{-# OPTIONS_GHC -fglasgow-exts #-}
{-# LANGUAGE GADTs, ScopedTypeVariables, EmptyDataDecls, FlexibleInstances, ImpredicativeTypes, NoMonoPatBinds #-}
-- XXX Despite what I think should be enough LANGUAGE options I still need -fglasgow-exts.
module Language.CMonad.Prim(E', E, V, runE, embed, auto, arrayU, liftArray, (=:), RValue) where
import Control.Monad
import Data.Array
import Data.Array.MArray

import Language.CMonad.MonadRef

-- |Generic value type, both l-values and r-values.
data E' v m a where
    E :: m a -> E' RValue m a              -- compound expressions, only r-values
    V :: m a -> (a -> m ()) -> E' v m a    -- variables, l-value or r-value

data LValue -- ^l-value tag
data RValue -- ^r-value tag

type E m a = E' RValue m a                 -- ^Type of r-values in monad /m/
type V m a = E' LValue m a                 -- ^Type of l-values in monad /m/

-- |Evaluate an expression to an expression in the corresponding monad.
{-# INLINE runE #-}
runE :: E' v m a -> m a
runE (E t) = t
runE (V t _) = t

-- |r-values form a monad.
instance (Monad m) => Monad (E' RValue m) where
    {-# INLINE return #-}
    return x = E $ return x
    {-# INLINE (>>=) #-}
    x >>= f = E $ do
        x' <- runE x
        runE (f x')

-- |Any expression in the underlying monad can be lifted to a C expression.
{-# INLINE embed #-}
embed :: m a -> E m a
embed = E

-- |A variable with a initial value.
{-# INLINE auto #-}
auto :: (MonadRef m r) => E m a -> E m (forall v . E' v m a)
auto x = E (do
    x' <- runE x
    r  <- newRef x'
    return (V (readRef r) (writeRef r))
  )

{-# INLINE liftArray #-}
liftArray :: forall arr m a i . (Ix i, MArray arr a m) =>
             arr i a -> E m (forall v . [E m i] -> E' v m a)
liftArray a = E ( do
    let ix :: [E m i] -> m i
        ix [i] = runE i
	{-# INLINE f #-}
	f is = V (ix is >>= readArray a) (\ x -> ix is >>= \ i -> writeArray a i x)
    return f
  )

-- |A un-initialized multi-dimensional array.  E.g., @arrayU [2,3]@ is a 2x3 array.
arrayU :: forall arr m a i . (Ix i, Num i, MArray arr a m, TheArray m arr) =>
       [E m i] -> E m (forall v . [E m i] -> E' v m a)
arrayU ss = E ( do
    ss' <- mapM runE ss
    let sz = product ss'
        ix :: [E m i] -> m i
        ix is = do
                    is' <- mapM runE is
                    when (length is' /= length ss') $
                        error "wrong number of indicies"
                    return $ foldr (\ (i, s) r -> r * s + i) 0 (zip is' ss')
    a <- newArray (0, product ss' - 1) undefined :: m (arr i a)
    return (\ is -> V (ix is >>= readArray a)
                      (\ x -> ix is >>= \ i -> writeArray a i x))
  )

-- |An C array initialized with a normal array.
arrayA :: forall arr m a i . (Ix i, MArray arr a m, TheArray m arr) =>
       Array i a -> E m (forall v . [E m i] -> E' v m a)
arrayA aa = E ( do
    a <- thaw aa :: m (arr i a)
    runE (liftArray a)
  )

-- |Assignment operator.
infix 0 =:
{-# INLINE (=:) #-}
(=:) :: (Monad m) => V m a -> E m a -> E m a
V _ asg =: e = do
    e' <- e
    E (asg e')
    return e'