{-# LANGUAGE TypeFamilies #-}
{- |
Maybe datatype implemented in continuation passing style.
-}
module LLVM.Extra.MaybeContinuation where

import qualified LLVM.Extra.Control as C
import LLVM.Extra.Control (ifThenElse, )

import LLVM.Extra.Class (Undefined, undefTuple, )
import qualified LLVM.Extra.Arithmetic as A
import LLVM.Core as LLVM
import LLVM.Util.Loop (Phi, ) -- (phis, addPhis, )

import Control.Monad.IO.Class (MonadIO(liftIO), )
import qualified Control.Applicative as App
import qualified Control.Monad as M

import Control.Monad.HT ((<=<), )
import Data.Tuple.HT (mapSnd, )

import Prelude hiding (and, iterate, map, zip, zipWith, writeFile, )


{- |
Isomorphic to @ReaderT (CodeGenFunction r z) (ContT z (CodeGenFunction r)) a@,
where the reader provides the block for 'Nothing'
and the continuation part manages the 'Just'.
-}
newtype T r z a =
   Cons {resolve ::
      CodeGenFunction r z ->
      (a -> CodeGenFunction r z) ->
      CodeGenFunction r z
   }


map :: (a -> CodeGenFunction r b) -> T r z a -> T r z b
map f (Cons m) = Cons $ \n j ->
   m n (j <=< f)

instance Functor (T r z) where
   fmap f (Cons m) = Cons $ \n j -> m n (j . f)

instance App.Applicative (T r z) where
   pure = return
   (<*>) = M.ap

instance Monad (T r z) where
   return a = lift (return a)
   (>>=) = bind

instance MonadIO (T r z) where
   liftIO = lift . liftIO

{- |
counterpart to Data.Maybe.HT.toMaybe
-}
withBool ::
   (Phi z) =>
   Value Bool -> CodeGenFunction r a -> T r z a
withBool b a =
   guard b >> lift a
{-
withBool b a = Cons $ \n j ->
   ifThenElse b (j =<< a) n
-}

fromBool ::
   (Phi z) =>
   CodeGenFunction r (Value Bool, a) ->
   T r z a
fromBool m = do
   (b,a) <- lift m
   guard b
   return a

toBool ::
   (Undefined a) =>
   T r (Value Bool, a) a -> CodeGenFunction r (Value Bool, a)
toBool (Cons m) =
   m (return (valueOf False, undefTuple)) (return . (,) (valueOf True))

isJust ::
   T r (Value Bool) a -> CodeGenFunction r (Value Bool)
isJust (Cons m) =
   m (return (valueOf False)) (const $ return (valueOf True))

lift :: CodeGenFunction r a -> T r z a
lift a = Cons $ \ _n j -> j =<< a

guard ::
   (Phi z) =>
   Value Bool -> T r z ()
guard b = Cons $ \n j ->
   ifThenElse b (j ()) n

{-
just :: CodeGenFunction r a -> T r z a
just a = Cons $ \ _n j -> j =<< a

nothing :: T r z a
nothing = Cons \n _j -> n
-}

bind ::
   T r z a ->
   (a -> T r z b) ->
   T r z b
bind (Cons ma) mb = Cons $ \n j ->
   ma n (\a -> resolve (mb a) n j)

{- |
If the returned position is smaller than the array size,
then returned final state is undefined.
-}
arrayLoop ::
   (Phi s, Undefined s, IsType a,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet i, CmpResult i ~ Bool) =>
   Value i ->
   Value (Ptr a) -> s ->
   (Value (Ptr a) -> s -> T r (Value Bool, s) s) ->
   CodeGenFunction r (Value i, s)
arrayLoop len ptr start loopBody =
   C.arrayLoopWithExit len ptr start $ \ptri s0 ->
      toBool (loopBody ptri s0)

arrayLoop2 ::
   (Phi s, Undefined s, IsType a, IsType b,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet i, CmpResult i ~ Bool) =>
   Value i ->
   Value (Ptr a) -> Value (Ptr b) -> s ->
   (Value (Ptr a) -> Value (Ptr b) -> s ->
      T r (Value Bool, (Value (Ptr b), s)) s) ->
   CodeGenFunction r (Value i, s)
arrayLoop2 len ptrA ptrB start loopBody =
   fmap (mapSnd snd) $
   arrayLoop len ptrA (ptrB,start) $ \ptrAi (ptrB0,s0) -> do
      s1 <- loopBody ptrAi ptrB0 s0
      ptrB1 <- lift $ A.advanceArrayElementPtr ptrB0
      return (ptrB1,s1)


fixedLengthLoop ::
   (Phi s, Undefined s,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet i, CmpResult i ~ Bool) =>
   Value i -> s ->
   (s -> T r (Value Bool, (Value i, s)) s) ->
   CodeGenFunction r (Value i, s)
fixedLengthLoop len start loopBody = do
   (_,(lastI,lastS)) <-
      C.whileLoopShared (valueOf True, (len, start)) $ \(cont,(i,s)) ->
         (A.and cont =<< A.cmp LLVM.CmpGT i (value LLVM.zero),
          resolve (loopBody s)
             (return (valueOf False, undefTuple))
             (\newS -> do
                newI <- A.dec i
                return (valueOf True, (newI, newS))))
   fmap (flip (,) lastS) $ A.sub len lastI


{-
In case of early exit we would not have a final state.
However, the loop could be in the T monad
and we could just propagate a Nothing.

whileLoop ::
   Phi a =>
   a ->
   (a -> T r z a) ->
   CodeGenFunction r a
whileLoop start check body = do
   top <- getCurrentBasicBlock
   loop <- newBasicBlock
   cont <- newBasicBlock
   exit <- newBasicBlock
   br loop

   defineBasicBlock loop
   state <- phis top start
   b <- check state
   condBr b cont exit
   defineBasicBlock cont
   res <- body state
   cont' <- getCurrentBasicBlock
   addPhis cont' state res
   br loop

   defineBasicBlock exit
   return state
-}