{-# LANGUAGE TypeOperators
            ,ScopedTypeVariables
            ,PostfixOperators
            ,NoMonomorphismRestriction
            ,BangPatterns
  #-}
{-# OPTIONS_GHC -fno-cse #-}
module Data.OI (
  OI
 ,(:->)
 ,iooi
 ,run
 ,(=:)
 ,(?)
 ,(#)
 ,idA
 ,(<.>)
 ,arrA
 ,firstA
 ,deTuple
 ,deList
 ,(<|)
 ,mapOI
 ,mapOI'
 ,zipWithOI
 ,zipWithOI'
 ,sequenceOI
 ,sequenceOI'
 ) where

import Control.Applicative
import Control.Concurrent.MVar
import Control.Parallel
import Control.Comonad
import System.IO.Unsafe

data OI a = OI { variable :: LeftValueOf (IO a), value :: a }
type a :-> b = OI a -> b

(=:)         ::  a -> a :-> a
x =: OI v y  =   assign ix v `pseq` y
                 where ix = return x

(?)           ::  a :-> a
(?) (OI _ x)  =   x

(#)    ::  a -> OI a
(#) x  =   OI { variable = reference (return x), value = x }

--

instance Functor OI where
  fmap f x = (#) (f (x?))

instance Monad OI where
  return = (#)
  x >>= f = f (x?)

instance Applicative OI where
  pure = (#)
  f <*> x = (#)((f?)(x?))

instance Extend OI where
  duplicate = (#)

instance Comonad OI where
  extract = (?)

--

idA :: a :-> a
idA = (?)

(<.>) :: (b :-> c) -> (a :-> b) -> (a :-> c)
f <.> g = f . (#) . g

--

arrA :: (a -> b) -> (a :-> b)
arrA f = f . (?)

firstA :: (a :-> b) -> (a,c) :-> (b,c)
firstA f ac = case deTuple ac of (x,z) -> (f x, (z?))

--

deTuple :: (a,b) :-> (OI a,OI b)
deTuple (OI vxy ~(x,y)) = assign io vxy `pseq` (OI vx x, OI vy y)
 where
  vx     = new ()
  vy     = new ()
  io     = (,) <$> unsafeInterleaveIO (dereference vx) <*> unsafeInterleaveIO (dereference vy)

deList :: [a] :-> Maybe (OI a, OI [a])
deList (OI vxxs xxs)
 = assign io vxxs
   `pseq` case xxs of
      x:xs -> Just (OI vx x, OI vxs xs)
      _    -> Nothing
   where
     vx  = new ()
     vxs = new ()
     io  = (:) <$> (unsafeInterleaveIO (dereference vx)) <*> (unsafeInterleaveIO (dereference vxs))

--
infixr 1 <|

(<|)  ::  (b -> c :-> d) -> (a :-> b) -> (a,c) :-> d
(f <| g) ac = case deTuple ac of (a,c) -> f (g a) c

mapOI :: (a :-> b) -> [a] :-> [b]
mapOI f xxs = case deList xxs of
  Just (x,xs) -> f x : mapOI f xs
  _           -> []

mapOI' :: (a :-> b) -> [a] :-> (OI [a],[b])
mapOI' f xxs = case deList xxs of
  Just (x,xs) -> (zs, f x:ys)
                 where (zs,ys) = mapOI' f xs
  _           -> (xxs,[])

zipWithOI :: (a -> b :-> c) -> [a] -> [b] :-> [c]
zipWithOI _ [] _  = []
zipWithOI f (x:xs) yys = case deList yys of
  Just (y,ys) -> f x y : zipWithOI f xs ys
  _           -> []

zipWithOI' :: (a -> b :-> c) -> [a] -> [b] :-> (OI [b],[c])
zipWithOI' _ []     yys = (yys,[])
zipWithOI' f (x:xs) yys = case deList yys of
  Just (y,ys) -> case zipWithOI' f xs ys of ~(rs,zs) -> (rs,f x y : zs)
  _           -> (yys,[])

sequenceOI :: [a :-> b] -> [a] :-> ()
sequenceOI (f:fs) xxs = case deList xxs of
  Just (x,xs) -> f x `pseq` sequenceOI fs xs
  Nothing     ->  ()
sequenceOI [] _ = ()

sequenceOI' :: [a :-> b] -> [a] :-> OI [a]
sequenceOI' (f:fs) xxs = case deList xxs of
  Just (x,xs) -> f x `pseq` sequenceOI' fs xs
  Nothing     ->  xxs
sequenceOI' [] xxs = xxs

--

iooi            ::  IO a -> (a :-> a)
iooi io (OI vix x)  =   assign io vix `pseq` x

run        ::  (a :-> b) -> IO b
run pmain  =   do { vx <- unsafeInterleaveIO newEmptyMVar
                  ; x  <- unsafeInterleaveIO (dereference vx)
                  ; return $! pmain (OI vx x)
                  }

--

type LeftValueOf = MVar

{-# INLINE new #-}
new :: () -> LeftValueOf a
new _ = unsafePerformIO $ newEmptyMVar

{-# INLINE reference #-}
reference :: a -> LeftValueOf a
reference = unsafePerformIO . newMVar

{-# INLINE dereference #-}
dereference :: LeftValueOf a -> a
dereference = unsafePerformIO . readMVar

{-# INLINE assign #-}
assign :: a -> LeftValueOf a -> a
assign !x v = unsafePerformIO 
            $ do { s <- tryPutMVar v x
                 ; if s then return x else readMVar v
                 }