{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
module Synthesizer.LLVM.Simple.SignalPacked where

import Synthesizer.LLVM.Simple.SignalPrivate (Core(Core), alter, )
import qualified Synthesizer.LLVM.Simple.Signal as Sig
import qualified Synthesizer.LLVM.Frame.SerialVector as Serial

import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.MaybeContinuation as Maybe
import qualified LLVM.Extra.Control as U
import qualified LLVM.Extra.Arithmetic as A
import LLVM.Extra.Class (undefTuple, )

import qualified LLVM.Core as LLVM
import LLVM.Core (valueOf, )

import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import Control.Monad (replicateM, )

import Data.Word (Word32, )

import NumericPrelude.Numeric as NP
import NumericPrelude.Base hiding (and, iterate, map, zip, zipWith, )



{- |
Convert a signal of scalar values into one using processor vectors.
If the signal length is not divisible by the chunk size,
then the last chunk is dropped.
-}
pack, packRotate ::
   (Sig.C signal, Serial.C v, a ~ Serial.Element v) =>
   signal a -> signal v
pack = packRotate

packRotate = alter (\(Core next start stop) -> Core
   (\param s -> do
      wInit <- Maybe.lift $ Serial.writeStart
      (w2,_,s2) <-
         Maybe.fromBool $
         U.whileLoop
            (valueOf True,
             (wInit,
              valueOf $ (fromIntegral $ Serial.sizeOfIterator wInit :: Word32),
              s))
            (\(cont,(_w0,i0,_s0)) ->
               A.and cont =<<
                  A.cmp LLVM.CmpGT i0 A.zero)
            (\(_,(w0,i0,s0)) -> Maybe.toBool $ do
               (a,s1) <- next param s0
               Maybe.lift $ do
                  w1 <- Serial.writeNext a w0
                  i1 <- A.dec i0
                  return (w1,i1,s1))
      v <- Maybe.lift $ Serial.writeStop w2
      return (v, s2))
   start
   stop)

{-
We could reformulate it in terms of WriteIterator
that accesses elements using LLVM.extract.
We might move the loop counter into the Iterator,
but we have to assert that the counter is not duplicated.

packIndex ::
   (Sig.C signal, Serial.C v, a ~ Serial.Element v) =>
   signal a -> signal v
packIndex = alter (\(Core next start stop) -> Core
   (\param s -> do
      (v2,_,s2) <-
         Maybe.fromBool $
         U.whileLoop
            (valueOf True, (undefTuple, A.zero, s))
            (\(cont,(v0,i0,_s0)) ->
               A.and cont =<<
                  A.cmp LLVM.CmpLT i0
                     (valueOf $ fromIntegral $ Serial.size v0))
            (\(_,(v0,i0,s0)) -> Maybe.toBool $ do
               (a,s1) <- next param s0
               Maybe.lift $ do
                  v1 <- Vector.insert i0 a v0
                  i1 <- A.inc i0
                  return (v1,i1,s1))
      return (v2, s2))
   start
   stop)
-}


{- |
Like 'pack' but duplicates the code for creating elements.
That is, for vectors of size n, the code of the input signal
will be emitted n times.
This is efficient only for simple input generators.
-}
packSmall ::
   (Sig.C signal, Serial.C v, a ~ Serial.Element v) =>
   signal a -> signal v
packSmall = alter (\(Core next start stop) -> Core
   (\param ->
      MS.runStateT $
      Serial.withSize $ \n ->
         MT.lift . Maybe.lift . Serial.assemble
         =<<
         replicateM n (MS.StateT $ next param))
   start
   stop)


unpack, unpackRotate ::
   (Sig.C signal,
    Serial.Read v, a ~ Serial.Element v, Serial.ReadIt v ~ itv, Memory.C itv) =>
   signal v -> signal a
unpack = unpackRotate

unpackRotate = alter (\(Core next start stop) -> Core
   (\context (i0,r0,s0) -> do
      endOfVector <-
         Maybe.lift $ A.cmp LLVM.CmpEQ i0 (valueOf (0::Word32))
      (i2,r2,s2) <-
         Maybe.fromBool $
         U.ifThen endOfVector (valueOf True, (i0,r0,s0)) $ do
            (cont1, (v1,s1)) <- Maybe.toBool $ next context s0
            r1 <- Serial.readStart v1
            return (cont1, (valueOf $ fromIntegral $ Serial.size v1, r1, s1))
      Maybe.lift $ do
         (a,r3) <- Serial.readNext r2
         i3 <- A.dec i2
         return (a, (i3,r3,s2)))
   (fmap (\s -> (A.zero, undefTuple, s)) . start)
   (\(_,_,state) -> stop state))


{-
We could reformulate it in terms of ReadIterator
that accesses elements using LLVM.extract.
We might move the loop counter into the Iterator,
but we have to assert that the counter is not duplicated.

unpackIndex ::
   (Serial.C v, a ~ Serial.Element v, Memory.C v) =>
   signal v -> signal a
unpackIndex = alter (\(Core next start stop) -> Core
   (\param (i0,v0,s0) -> do
      endOfVector <-
         Maybe.lift $ A.cmp LLVM.CmpGE i0
            (valueOf $ fromIntegral $ Serial.size v0)
      (i2,v2,s2) <-
         Maybe.fromBool $
         U.ifThen endOfVector (valueOf True, (i0,v0,s0)) $ do
            (cont1, (v1,s1)) <- Maybe.toBool $ next param s0
            return (cont1, (A.zero, v1, s1))
      Maybe.lift $ do
         a <- Vector.extract i2 v2
         i3 <- A.inc i2
         return (a, (i3,v2,s2)))
   (\p -> do
      s <- start p
      let v = undefTuple
      return (valueOf $ fromIntegral $ Serial.size v, v, s))
   stop)
-}