{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{- |
Maybe datatype implemented in continuation passing style.
-}
module LLVM.Extra.MaybeContinuation where

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

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

import qualified Control.Applicative as App
import qualified Control.Monad as M

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

import Prelude hiding (fmap, and, iterate, map, zip, zipWith, writeFile, )
import qualified Prelude as P


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

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

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, IsType a,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet 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 =
   U.arrayLoopWithExit len ptr start $ \ptri s0 ->
      toBool (loopBody ptri s0)

{-
arrayLoop len ptr start loopBody = do
   top <- getCurrentBasicBlock
   loop <- newBasicBlock
   body <- newBasicBlock
   exit <- newBasicBlock

   br loop

   defineBasicBlock loop
   i <- phi [(len, top)]
   p <- phi [(ptr, top)]
   vars <- phis top start
   t <- A.icmp IntNE i (value LLVM.zero)
   condBr t body exit

   defineBasicBlock body
   loopBody p vars
      (br exit)
      (\vars' -> do
         next <- getCurrentBasicBlock
         addPhis next vars vars'

         i' <- A.dec i
         p' <- A.advanceArrayElementPtr p

         addPhiInputs i [(i', next)]
         addPhiInputs p [(p', next)]
         br loop)

   defineBasicBlock exit
   pos <- sub len i
   return (pos, vars)
-}

arrayLoop2 ::
   (Phi s, IsType a, IsType b,
    Num i, IsConst i, IsInteger i, IsFirstClass i, CmpRet 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 =
   P.fmap (mapSnd snd) $
   arrayLoop len ptrA (ptrB,start) $ \ptrAi (ptrBi,s0) -> do
      s1 <- loopBody ptrAi ptrBi s0
      ptrBi' <- lift $ A.advanceArrayElementPtr ptrBi
      return (ptrBi',s1)

{-
a specialised variant of whileLoop might also be useful
-}