ProcessPrivate.hs

{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.CausalParameterized.ProcessPrivate where

import qualified Synthesizer.LLVM.Parameterized.SignalPrivate as Sig
import qualified Synthesizer.LLVM.Parameter as Param
import qualified Synthesizer.LLVM.Causal.Process as Causal
import Synthesizer.LLVM.Causal.ProcessPrivate (loopNext, )

import qualified Synthesizer.Causal.Class as CausalClass

import qualified LLVM.Extra.Control as C
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.MaybeContinuation as MaybeCont
import qualified LLVM.Extra.Memory as Memory
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, Undefined, )

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

import Types.Data.Num (d1, )

import qualified Control.Monad.HT as M
import qualified Control.Arrow    as Arr
import qualified Control.Category as Cat
import Control.Arrow (arr, (^<<), (<<<), (&&&), )
import Control.Applicative (Applicative, pure, (<*>), )
import Data.Tuple.HT (mapSnd, )

import qualified Synthesizer.LLVM.Storable.Vector as SVU
import qualified Data.StorableVector as SV

import Data.Word (Word32, )
import Foreign.Storable.Tuple ()
import Foreign.Storable (Storable, )
import Foreign.ForeignPtr (ForeignPtr, touchForeignPtr, )
import Foreign.Ptr (Ptr, )

import qualified Number.Ratio as Ratio
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

import NumericPrelude.Numeric
import NumericPrelude.Base hiding (and, iterate, map, zip, zipWith, take, takeWhile, init, )

import qualified Prelude as P


data T p a b =
   forall context state ioContext parameters.
      (Storable parameters,
       MakeValueTuple parameters,
       Memory.C (ValueTuple parameters),
       Memory.C context,
       Memory.C state) =>
   Cons
      (forall r c.
       (Phi c) =>
       context ->
       a -> state -> MaybeCont.T r c (b, state))
          -- compute next value
      (forall r.
       ValueTuple parameters ->
       CodeGenFunction r (context, state))
          -- initial state
      (forall r.
       context -> state ->
       CodeGenFunction r ())
          -- cleanup
      (p -> IO (ioContext, parameters))
          {- initialization from IO monad
          This will be run within Unsafe.performIO,
          so no observable In/Out actions please!
          -}
      (ioContext -> IO ())
          -- finalization from IO monad, also run within Unsafe.performIO

instance CausalClass.C (T p) where
   type SignalOf (T p) = Sig.T p
   type ProcessOf (Sig.T p) = T p
   toSignal = toSignal
   fromSignal = fromSignal

instance Causal.C (T p) where
   simple next start =
      simple (\() -> next) (\() -> fmap ((,) ()) start) (pure ())

   loopConst init (Cons next start stop createIOContext deleteIOContext) =
      Cons
         (loopNext next)
         (fmap (mapSnd ((,) init)) . start)
         (loopStop stop)
         createIOContext
         deleteIOContext

   replicateControlled n = replicateControlled $ pure n


simple ::
   (Storable parameters,
    MakeValueTuple parameters, ValueTuple parameters ~ paramValue,
    Memory.C paramValue,
    Memory.C context,
    Memory.C state) =>
   (forall r c.
    (Phi c) =>
    context -> a -> state -> MaybeCont.T r c (b, state)) ->
   (forall r.
    paramValue ->
    CodeGenFunction r (context, state)) ->
   Param.T p parameters -> T p a b
simple f start param =
   Param.with param $ \get value -> Cons f
      (start . value)
      (const $ const $ return ())
      (return . (,) () . get)
      (const $ return ())


toSignal :: T p () a -> Sig.T p a
toSignal (Cons next start stop createIOContext deleteIOContext) = Sig.Cons
   (\p -> next p ())
   start stop
   createIOContext deleteIOContext

fromSignal :: Sig.T p b -> T p a b
fromSignal (Sig.Cons next start stop createIOContext deleteIOContext) = Cons
   (\p _ -> next p)
   start stop
   createIOContext deleteIOContext


mapAccum ::
   (Storable pnh, MakeValueTuple pnh, ValueTuple pnh ~ pnl, Memory.C pnl,
    Storable psh, MakeValueTuple psh, ValueTuple psh ~ psl, Memory.C psl,
    Memory.C s) =>
   (forall r. pnl -> a -> s -> CodeGenFunction r (b,s)) ->
   (forall r. psl -> CodeGenFunction r s) ->
   Param.T p pnh ->
   Param.T p psh ->
   T p a b
mapAccum next start selectParamN selectParamS =
   simple
      (\p a s -> MaybeCont.lift $ next p a s)
      (\(n,s) -> fmap ((,) n) $ start s)
      (selectParamN &&& selectParamS)


map ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl) =>
   (forall r. pl -> a -> CodeGenFunction r b) ->
   Param.T p ph ->
   T p a b
map f selectParamF =
   mapAccum
      (\p a s -> fmap (flip (,) s) $ f p a)
      (const $ return ())
      selectParamF
      (return ())

mapSimple ::
   (forall r. a -> CodeGenFunction r b) ->
   T p a b
mapSimple f =
   map (const f) (return ())

zipWith ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl) =>
   (forall r. pl -> a -> b -> CodeGenFunction r c) ->
   Param.T p ph ->
   T p (a,b) c
zipWith f =
   map (uncurry . f)

zipWithSimple ::
   (forall r. a -> b -> CodeGenFunction r c) ->
   T p (a,b) c
zipWithSimple f =
   mapSimple (uncurry f)


apply :: T p a b -> Sig.T p a -> Sig.T p b
apply = CausalClass.apply

feedFst :: Sig.T p a -> T p b (a,b)
feedFst = CausalClass.feedFst

feedSnd :: Sig.T p a -> T p b (b,a)
feedSnd = CausalClass.feedSnd


{-
Very similar to 'apply',
since 'apply' can be considered being of type
@T p a b -> T p () a -> T p () b@.
-}
compose :: T p a b -> T p b c -> T p a c
compose
      (Cons nextA startA stopA createIOContextA deleteIOContextA)
      (Cons nextB startB stopB createIOContextB deleteIOContextB) =
   Cons
      (composeNext MaybeCont.onFail stopA stopB nextA nextB)
      (composeStart startA startB)
      (composeStop stopA stopB)
      (composeCreate createIOContextA createIOContextB)
      (composeDelete deleteIOContextA deleteIOContextB)

composeNext ::
   (Monad maybe) =>
   (forall x. code () -> maybe x -> maybe x) ->
   (contextA -> stateA -> code ()) ->
   (contextB -> stateB -> code ()) ->
   (contextA -> a -> stateA -> maybe (b, stateA)) ->
   (contextB -> b -> stateB -> maybe (c, stateB)) ->
   (contextA, contextB) ->
   a ->
   (stateA, stateB) ->
   maybe (c, (stateA, stateB))
composeNext onFail stopA stopB nextA nextB (paramA, paramB) a (sa0,sb0) = do
   (b,sa1) <-
      onFail (stopB paramB sb0) $
      nextA paramA a sa0
   (c,sb1) <-
      onFail (stopA paramA sa1) $
      nextB paramB b sb0
   return (c, (sa1,sb1))

composeStart ::
   Monad m =>
   (paramA -> m (contextA, stateA)) ->
   (paramB -> m (contextB, stateB)) ->
   (paramA, paramB) -> m ((contextA, contextB), (stateA, stateB))
composeStart = Sig.combineStart

composeStop ::
   Monad m =>
   (contextA -> stateA -> m ()) ->
   (contextB -> stateB -> m ()) ->
   (contextA, contextB) -> (stateA, stateB) -> m ()
composeStop = Sig.combineStop

composeCreate ::
   Monad m =>
   (p -> m (ioContextA, contextA)) ->
   (p -> m (ioContextB, contextB)) ->
   p -> m ((ioContextA, ioContextB), (contextA, contextB))
composeCreate = Sig.combineCreate

composeDelete ::
   (Monad m) =>
   (ca -> m ()) -> (cb -> m ()) -> (ca, cb) -> m ()
composeDelete = Sig.combineDelete


first :: T p b c -> T p (b, d) (c, d)
first (Cons next start stop createIOContext deleteIOContext) = Cons
   (firstNext next) start stop
   createIOContext deleteIOContext

firstNext ::
   Monad m =>
   (context -> a -> s -> m (b, s)) ->
   context -> (a, c) -> s -> m ((b, c), s)
firstNext next context (b,d) sa0 = do
   (c,sa1) <- next context b sa0
   return ((c,d), sa1)


{- |
serial replication

But you may also use it for a parallel replication, see 'replicateParallel'.
-}
replicateControlled ::
   (Undefined x, Phi x) =>
   Param.T p Int -> T p (c,x) x -> T p (c,x) x
replicateControlled n (Cons next start stop createIOContext deleteIOContext) =
   case Param.word32 n of
      n32 -> Cons
         (\(len, cs) ->
            replicateControlledNext next stop (Param.value n32 len, cs))
         (\(len, param) ->
            replicateControlledStart start (Param.value n32 len, param))
         (\(len, cs) ->
            replicateControlledStop stop (Param.value n32 len, cs))
         (\p ->
            replicateControlledCreate $
               M.replicate (Param.get n p) (createIOContext p))
         (replicateControlledDelete deleteIOContext)

replicateControlledNext ::
   (Memory.C context, Memory.C state,
    contextState ~
       LLVM.Struct (Memory.Struct context, (Memory.Struct state, ())),
    Phi z, Phi a, Undefined a) =>
   (forall z0. (Phi z0) =>
    context -> (ctrl, a) -> state ->
    MaybeCont.T r z0 (a, state)) ->
   (context -> state -> CodeGenFunction r ()) ->
   (Value Word32, Value (Ptr contextState)) ->
   (ctrl, a) ->
   () ->
   MaybeCont.T r z (a, ())
replicateControlledNext next stop (len, contextStates) (c,a) () =
   MaybeCont.fromMaybe $ fmap (\(_,ms) -> fmap (flip (,) ()) ms) $
      MaybeCont.arrayLoop len contextStates a $
            \contextStatePtr a0 -> do
         (context, s0) <- MaybeCont.lift $ Memory.load contextStatePtr
         (a1,s1) <-
            MaybeCont.onFail
               (replicateControlledStopExcept
                  stop len contextStates contextStatePtr) $
            next context (c,a0) s0
         MaybeCont.lift $
            Memory.store s1 =<< LLVM.getElementPtr0 contextStatePtr (d1, ())
         return a1

replicateControlledStopExcept ::
   (Memory.C a, Memory.C b,
    ab ~ LLVM.Struct (Memory.Struct a, (Memory.Struct b, ()))) =>
   (a -> b -> CodeGenFunction r ()) ->
   Value Word32 ->
   Value (Ptr ab) ->
   Value (Ptr ab) ->
   CodeGenFunction r ()
replicateControlledStopExcept stop len contextStates contextStatePtr =
   C.arrayLoop len contextStates () $ \ptr () -> do
      b <- A.cmp LLVM.CmpNE ptr contextStatePtr
      C.ifThen b () $ uncurry stop =<< Memory.load ptr

_replicateControlledNext ::
   (Memory.C context, Memory.C state,
    contextState ~
       LLVM.Struct (Memory.Struct context, (Memory.Struct state, ())),
    Phi z, Phi a, Undefined a) =>
   (forall z0. (Phi z0) =>
    context -> (ctrl, a) -> state ->
    MaybeCont.T r z0 (a, state)) ->
   (Value Word32, Value (Ptr contextState)) ->
   (ctrl, a) ->
   () ->
   MaybeCont.T r z (a, ())
_replicateControlledNext next (len, contextStates) (c,a) () =
   fmap (flip (,) ()) $ MaybeCont.fromBool $ fmap snd $
   C.arrayLoopWithExit len contextStates (valueOf True, a) $
         \contextStatePtr (_,a0) -> do
      (context, s0) <- Memory.load contextStatePtr
      (cont, (a1,s1)) <- MaybeCont.toBool $ next context (c,a0) s0
      Memory.store s1 =<< LLVM.getElementPtr0 contextStatePtr (d1, ())
      return (cont, (cont,a1))

replicateControlledStart ::
   (Memory.C a, Memory.C b) =>
   (a -> CodeGenFunction r b) ->
   (Value Word32, Value (Ptr (Memory.Struct a))) ->
   CodeGenFunction r ((Value Word32, Value (Ptr (Memory.Struct b))), ())
replicateControlledStart start (len, params) = do
   contextStates <- LLVM.arrayMalloc len
   C.arrayLoop2 len params contextStates () $ \paramPtr statePtr () ->
      flip Memory.store statePtr =<< start =<< Memory.load paramPtr
   return ((len, contextStates), ())

replicateControlledStop ::
   (Memory.C a, Memory.C b,
    ab ~ LLVM.Struct (Memory.Struct a, (Memory.Struct b, ()))) =>
   (a -> b -> CodeGenFunction r ()) ->
   (Value Word32, Value (Ptr ab)) ->
   () ->
   CodeGenFunction r ()
replicateControlledStop stop (len, contextStates) () = do
   C.arrayLoop len contextStates () $ \contextStatePtr () ->
      uncurry stop =<< Memory.load contextStatePtr
   LLVM.free contextStates


replicateControlledCreate ::
   (Monad m, MakeValueTuple b, Storable b, ValueTuple b ~ bl, Memory.C bl) =>
   m [(a, b)] ->
   m (([a], ForeignPtr b),
      (Word32, Ptr (Memory.Struct bl)))
replicateControlledCreate createIOContexts = do
   (ioContexts, params) <- M.lift unzip createIOContexts
   let (fptr, ptr, len) = SVU.unsafeToPointers $ SV.pack params
   return
      ((ioContexts, fptr),
       (fromIntegral len, ptr))

replicateControlledDelete ::
   (a -> IO ()) ->
   ([a], ForeignPtr b) -> IO ()
replicateControlledDelete deleteIOContext (ioContexts, fptr) = do
   mapM_ deleteIOContext ioContexts
   touchForeignPtr fptr


instance Cat.Category (T p) where
   id = mapSimple return
   (.) = flip compose

instance Arr.Arrow (T p) where
   arr f = mapSimple (return . f)
   first = first


instance Functor (T p a) where
   fmap = (^<<)

instance Applicative (T p a) where
   pure x = Arr.arr (const x)
   f <*> x = uncurry ($) ^<< f&&&x


instance (A.Additive b) => Additive.C (T p a b) where
   zero = pure A.zero
   negate x = mapSimple A.neg <<< x
   x + y = zipWithSimple A.add <<< x&&&y
   x - y = zipWithSimple A.sub <<< x&&&y

instance (A.PseudoRing b, A.IntegerConstant b) => Ring.C (T p a b) where
   one = pure A.one
   fromInteger n = pure (A.fromInteger' n)
   x * y = zipWithSimple A.mul <<< x&&&y

instance (A.Field b, A.RationalConstant b) => Field.C (T p a b) where
   fromRational' x = pure (A.fromRational' $ Ratio.toRational98 x)
   x / y = zipWithSimple A.fdiv <<< x&&&y


instance (A.PseudoRing b, A.Real b, A.IntegerConstant b) => P.Num (T p a b) where
   fromInteger n = pure (A.fromInteger' n)
   negate x = mapSimple A.neg <<< x
   x + y = zipWithSimple A.add <<< x&&&y
   x - y = zipWithSimple A.sub <<< x&&&y
   x * y = zipWithSimple A.mul <<< x&&&y
   abs x = mapSimple A.abs <<< x
   signum x = mapSimple A.signum <<< x

instance (A.Field b, A.Real b, A.RationalConstant b) => P.Fractional (T p a b) where
   fromRational x = pure (A.fromRational' x)
   x / y = zipWithSimple A.fdiv <<< x&&&y


{- |
Not quite the loop of ArrowLoop
because we need a delay of one time step
and thus an initialization value.

For a real ArrowLoop.loop, that is a zero-delay loop,
we would formally need a MonadFix instance of CodeGenFunction.
But this will not become reality, since LLVM is not able to re-order code
in a way that allows to access a result before creating the input.
-}
loop ::
   (Storable c, MakeValueTuple c, ValueTuple c ~ cl, Memory.C cl) =>
   Param.T p c -> T p (a,cl) (b,cl) -> T p a b
loop initial (Cons next start stop createIOContext deleteIOContext) =
   Param.with initial $ \getInitial valueInitial -> Cons
      (loopNext next)
      (\(i,p) -> fmap (mapSnd ((,) (valueInitial i))) $ start p)
      (loopStop stop)
      (\p -> do
         (ctx, param) <- createIOContext p
         return (ctx, (getInitial p, param)))
      deleteIOContext

loopStop :: (context -> state -> m) -> context -> (c, state) -> m
loopStop stop ctx (_c,s) = stop ctx s


takeWhile ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl) =>
   (forall r. pl -> a -> CodeGenFunction r (Value Bool)) ->
   Param.T p ph ->
   T p a a
takeWhile check selectParam = simple
   (\p a () -> do
      MaybeCont.guard =<< MaybeCont.lift (check p a)
      return (a, ()))
   (\p -> return (p, ()))
   selectParam


take ::
   Param.T p Int ->
   T p a a
take len =
   snd ^<<
   takeWhile (const $ A.cmp LLVM.CmpLT A.zero . fst) (return ()) <<<
   feedFst
      (Sig.iterate (const A.dec) (return ())
         ((fromIntegral :: Int -> Word32) . max 0 ^<< len))


{- |
The first output value is the initial value.
Thus 'integrate' delays by one sample compared with 'integrateSync'.
-}
integrate ::
   (Storable a, A.Additive al,
    MakeValueTuple a, ValueTuple a ~ al, Memory.C al) =>
   Param.T p a ->
   T p al al
integrate =
   flip loop (arr snd &&& zipWithSimple A.add)

integrateSync ::
   (Storable a, A.Additive al,
    MakeValueTuple a, ValueTuple a ~ al, Memory.C al) =>
   Param.T p a ->
   T p al al
integrateSync =
   flip loop ((\a -> (a,a)) ^<< zipWithSimple A.add)