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

import qualified LLVM.Extra.Maybe as Maybe
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Control as C
import LLVM.Extra.Control (ifThenElse, )
import LLVM.Extra.Class (Undefined, undefTuple, )

import qualified LLVM.Core as LLVM
import LLVM.Core
   (Value, value, valueOf,
    CodeGenFunction,
    IsConst, IsType, IsFirstClass, IsInteger,
    CmpRet, CmpResult, )
import LLVM.Util.Loop (Phi, ) -- (phis, addPhis, )

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

import Foreign.Ptr (Ptr, )
import Data.Tuple.HT (mapSnd, )

import Prelude hiding (map, )


{- |
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))


fromPlainMaybe :: (Phi z) => Maybe.T a -> T r z a
fromPlainMaybe (Maybe.Cons b a) = guard b >> return a

fromMaybe :: (Phi z) => CodeGenFunction r (Maybe.T a) -> T r z a
fromMaybe m = lift m >>= fromPlainMaybe

toMaybe ::
   (Undefined a) =>
   T r (Maybe.T a) a -> CodeGenFunction r (Maybe.T a)
toMaybe (Cons m) =
   m (return Maybe.nothing) (return . Maybe.just)


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 :: 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)

{- |
Run an exception handler if the Maybe-action fails.
The exception is propagated.
That is, the handler is intended for a cleanup procedure.
-}
onFail :: CodeGenFunction r () -> T r z a -> T r z a
onFail handler m = Cons $ \n j -> resolve m (handler >> n) j

{- |
Run the first action and if that fails run the second action.
If both actions fail, then the composed action fails, too.
-}
alternative ::
   (Phi z, Undefined a) =>
   T r (Maybe.T a) a -> T r (Maybe.T a) a -> T r z a
alternative x y =
   fromMaybe $ resolve x (toMaybe y) (return . Maybe.just)


fixedLengthLoop ::
   (Phi s, Undefined s,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet i, CmpResult i ~ Bool) =>
   Value i -> s ->
   (s -> T r (Maybe.T s) s) ->
   CodeGenFunction r (Value i, Maybe.T s)
fixedLengthLoop len start loopBody = do
   (vars, i) <-
      C.loopWithExit (start, len)
         (\(s0, i) -> do
            counterRunning <- A.cmp LLVM.CmpGT i (value LLVM.zero)
            (running, ms1) <-
               C.ifThen counterRunning (valueOf False, Maybe.just s0) $
               fmap (\m -> (Maybe.isJust m, m)) $ toMaybe $ loopBody s0
            return (running, (ms1, i)))
         (\(ms, i) ->
            fmap ((,) (Maybe.fromJust ms)) $ A.dec i)
   pos <- A.sub len i
   return (pos, vars)


{- |
If the returned position is smaller than the array size,
then returned final state is 'Maybe.nothing'.
-}
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 (Maybe.T (Value (Ptr a), s)) s) ->
   CodeGenFunction r (Value i, Maybe.T s)
arrayLoop len ptr start loopBody =
   fmap (mapSnd (fmap snd)) $
   fixedLengthLoop len (ptr,start) $ \(ptr0,s0) -> do
      s1 <- loopBody ptr0 s0
      ptr1 <- lift $ A.advanceArrayElementPtr ptr0
      return (ptr1,s1)


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 (Maybe.T (Value (Ptr a), (Value (Ptr b), s))) s) ->
   CodeGenFunction r (Value i, Maybe.T s)
arrayLoop2 len ptrA ptrB start loopBody =
   fmap (mapSnd (fmap snd)) $
   arrayLoop len ptrA (ptrB,start) $ \ptrAi (ptrB0,s0) -> do
      s1 <- loopBody ptrAi ptrB0 s0
      ptrB1 <- lift $ A.advanceArrayElementPtr ptrB0
      return (ptrB1,s1)


{-
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
-}