Functional.hs

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.CausalParameterized.Functional (
   T,
   lift,
   ($&), (&|&),
   compile,
   ) where

import qualified Synthesizer.LLVM.CausalParameterized.Process as Causal

import qualified LLVM.Extra.MaybeContinuation as Maybe
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A

import qualified LLVM.Core as LLVM
import LLVM.Extra.Class (MakeValueTuple, )
import LLVM.Util.Loop (Phi, )
import LLVM.Core (CodeGenFunction, IsSized, )

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 qualified Control.Monad.Trans.State as State
import qualified Control.Monad.Trans.Class as Trans
import Control.Monad.Trans.State (StateT, )

import qualified Data.Vault as Vault
import Data.Vault (Vault, )
import qualified Control.Category as Cat
import Control.Arrow (Arrow, (>>^), (&&&), arr, first, )
import Control.Category (Category, (.), )
import Control.Monad (liftM2, )
import Control.Applicative (Applicative, (<*>), pure, )

import Foreign.Storable (Storable, )

import System.IO.Unsafe (unsafePerformIO, )

import Prelude hiding ((.), )


newtype T p inp out = Cons (Code p inp out)


-- | similar to @Causal.T p a b@
data Code p a b =
   forall state packed size ioContext
        startParamTuple startParamValue startParamPacked startParamSize
        nextParamTuple  nextParamValue  nextParamPacked  nextParamSize.
      (Storable startParamTuple,
       Storable nextParamTuple,
       MakeValueTuple startParamTuple startParamValue,
       MakeValueTuple nextParamTuple  nextParamValue,
       Memory.C startParamValue startParamPacked,
       Memory.C nextParamValue  nextParamPacked,
       LLVM.IsSized startParamPacked startParamSize,
       LLVM.IsSized nextParamPacked  nextParamSize,
       Memory.C state packed,
       IsSized packed size) =>
   Code
      (forall r c.
       (Phi c) =>
       nextParamValue ->
       a -> state ->
       StateT Vault (Maybe.T r c) (b, state))
          -- compute next value
      (forall r.
       startParamValue ->
       CodeGenFunction r state)
          -- initial state
      (p -> IO (ioContext, (nextParamTuple, startParamTuple)))
          {- initialization from IO monad
          This will be run within unsafePerformIO,
          so no observable In/Out actions please!
          -}
      (ioContext -> IO ())
          -- finalization from IO monad, also run within unsafePerformIO



instance Category (Code p) where
   id = arr id
   Code nextB startB createIOContextB deleteIOContextB .
      Code nextA startA createIOContextA deleteIOContextA = Code
         (\(paramA, paramB) a (sa0,sb0) ->
            do (b,sa1) <- nextA paramA a sa0
               (c,sb1) <- nextB paramB b sb0
               return (c, (sa1,sb1)))
         (\(paramA, paramB) ->
            liftM2 (,)
               (startA paramA)
               (startB paramB))
         (\p -> do
            (ca,(nextParamA,startParamA)) <- createIOContextA p
            (cb,(nextParamB,startParamB)) <- createIOContextB p
            return ((ca,cb),
               ((nextParamA,  nextParamB),
                (startParamA, startParamB))))
         (\(ca,cb) ->
            deleteIOContextA ca >>
            deleteIOContextB cb)

instance Arrow (Code p) where
   arr f = Code
      (\ _p a state -> return (f a, state))
      (const $ return ())
      (const $ return ((),((),())))
      (const $ return ())
   first (Code next start create delete) = Code
      (\ioContext (b,d) sa0 ->
         do (c,sa1) <- next ioContext b sa0
            return ((c,d), sa1))
      start create delete


{-
We must not define Category and Arrow instances
because in osci***osci the result of osci would be shared,
although it depends on the particular input.

instance Category (T p) where
   id = tagUnique Cat.id
   Cons a . Cons b = tagUnique (a . b)

instance Arrow (T p) where
   arr f = tagUnique $ arr f
   first (Cons a) = tagUnique $ first a
-}

instance Functor (T p inp) where
   fmap f (Cons x) =
      tagUnique $ x >>^ f

instance Applicative (T p inp) where
   pure a = tagUnique $ arr (const a)
   f <*> x = fmap (uncurry ($))  $  f &|& x


instance (A.Additive b) => Additive.C (T p a b) where
   zero = lift Additive.zero
   negate x = Causal.mapSimple A.neg $& x
   x + y = Causal.zipWithSimple A.add $& x&|&y
   x - y = Causal.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 = Causal.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 = Causal.zipWithSimple A.fdiv $& x&|&y


infixr 0 $&

($&) :: Causal.T p b c -> T p a b -> T p a c
f $& (Cons b) =
   tagUnique $  liftCode f . b


infixr 3 &|&

(&|&) :: T p a b -> T p a c -> T p a (b,c)
Cons b &|& Cons c =
   tagUnique $  b &&& c


liftCode :: Causal.T p inp out -> Code p inp out
liftCode (Causal.Cons next start create delete) =
   Code
      (\p a state -> Trans.lift (next p a state))
      start create delete

lift :: Causal.T p inp out -> T p inp out
lift = tagUnique . liftCode

tag :: Vault.Key out -> Code p inp out -> T p inp out
tag key (Code next start create delete) =
   Cons $
   Code
      (\p a s0 -> do
         mb <- State.gets (Vault.lookup key)
         case mb of
            Just b -> return (b,s0)
            Nothing -> do
               bs@(b,_) <- next p a s0
               State.modify (Vault.insert key b)
               return bs)
      start create delete

-- dummy for debugging
_tag :: Vault.Key out -> Code p inp out -> T p inp out
_tag _ = Cons

tagUnique :: Code p inp out -> T p inp out
tagUnique code =
   unsafePerformIO $
   fmap (flip tag code) Vault.newKey

initialize :: Code p inp out -> Causal.T p inp out
initialize (Code next start create delete) =
   Causal.Cons
      (\p a state -> State.evalStateT (next p a state) Vault.empty)
      start create delete

compile :: T p inp out -> Causal.T p inp out
compile (Cons code) = initialize code