{-# LANGUAGE TypeFamilies #-}
module LLVM.Extra.MaybePrivate where

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

import qualified LLVM.Core as LLVM
import LLVM.Core (Value, valueOf, CodeGenFunction, )
import LLVM.Util.Loop (Phi, phis, addPhis, )

import qualified Control.Monad as Monad

import Prelude hiding (map, sequence)


{- |
If @isJust = False@, then @fromJust@ is an @undefTuple@.
-}
data T a = Cons {isJust :: Value Bool, fromJust :: a}


instance Functor T where
   fmap f (Cons b a) = Cons b (f a)

instance (Phi a) => Phi (T a) where
   phis bb (Cons b a) = Monad.liftM2 Cons (phis bb b) (phis bb a)
   addPhis bb (Cons b0 a0) (Cons b1 a1) =
      addPhis bb b0 b1 >> addPhis bb a0 a1


{- |
counterpart to 'maybe'
-}
run ::
   (Phi b) =>
   T a ->
   CodeGenFunction r b ->
   (a -> CodeGenFunction r b) ->
   CodeGenFunction r b
run (Cons b a) n j =
   ifThenElse b (j a) n

for ::
   T a ->
   (a -> CodeGenFunction r ()) ->
   CodeGenFunction r ()
for = flip run (return ())

{- |
counterpart to 'Data.Maybe.fromMaybe' with swapped arguments
-}
select ::
   (C.Select a) =>
   T a ->
   a ->
   CodeGenFunction r a
select (Cons b a) d = C.select b a d

alternative ::
   (C.Select a) =>
   T a -> T a -> CodeGenFunction r (T a)
alternative (Cons b0 a0) (Cons b1 a1) =
   Monad.liftM2 Cons
      (LLVM.or b0 b1)
      (C.select b0 a0 a1)


{- |
counterpart to Data.Maybe.HT.toMaybe
-}
fromBool :: Value Bool -> a -> T a
fromBool = Cons

toBool :: T a -> (Value Bool, a)
toBool (Cons b a) = (b,a)

just :: a -> T a
just = Cons (valueOf True)

nothing :: a -> T a
nothing undef = Cons (valueOf False) undef

getIsNothing :: T a -> CodeGenFunction r (Value Bool)
getIsNothing (Cons b _a) = LLVM.inv b


lift2 ::
   (a -> b -> c) ->
   T a -> T b -> CodeGenFunction r (T c)
lift2 f (Cons b0 a0) (Cons b1 a1) =
   Monad.liftM (flip Cons (f a0 a1)) (LLVM.and b0 b1)

sequence :: T (CodeGenFunction r a) -> CodeGenFunction r (T a)
sequence (Cons b0 a0) =
   Monad.liftM (Cons b0) a0

traverse ::
   (a -> CodeGenFunction r b) ->
   T a -> CodeGenFunction r (T b)
traverse f = sequence . fmap f

liftM2 ::
   (a -> b -> CodeGenFunction r c) ->
   T a -> T b -> CodeGenFunction r (T c)
liftM2 f ma mb = Monad.join $ fmap sequence $ lift2 f ma mb


maybeArg ::
   (Phi b) =>
   b ->
   (a -> CodeGenFunction r (T b)) ->
   T a -> CodeGenFunction r (T b)
maybeArg undef f m = run m (return $ nothing undef) f