{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{- |
Represent a vector of Stereo values in two vectors
that store the values in an interleaved way.
That is:

> vector0[0] = left[0]
> vector0[1] = right[0]
> vector0[2] = left[1]
> vector0[3] = right[1]
> vector1[0] = left[2]
> vector1[1] = right[2]
> vector1[2] = left[3]
> vector1[3] = right[3]

This representation is not very useful for computation,
but necessary as intermediate representation for interfacing with memory.
SSE/SSE2 have the instructions UNPACK(L|H)P(S|D) that interleave efficiently.
-}
module Synthesizer.LLVM.Frame.StereoInterleaved (
   T,
   Value,
   interleave,
   deinterleave,
   ) where

import qualified Synthesizer.LLVM.Frame.Stereo as Stereo

import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Class as Class
import qualified LLVM.Core as LLVM
import LLVM.Extra.Class
   (Undefined, undefTuple,
    MakeValueTuple, valueTupleOf, )
import LLVM.Core
   (Vector, Struct, IsSized, )
import LLVM.Util.Loop (Phi, phis, addPhis, )

import qualified LLVM.Extra.Memory as Memory
-- import qualified LLVM.Extra.Control as C
-- import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Vector as Vector
import qualified Data.TypeLevel.Num as TypeNum

import Foreign.Ptr (castPtr, )
import qualified Foreign.Storable as St
-- import Data.Word (Word32, )

import Control.Monad (liftM2, )
import Control.Applicative (liftA2, )

import Data.Tuple.HT (mapPair, )


data T n a = Cons (Vector n a) (Vector n a)

data Value n a = Value (LLVM.Value (Vector n a)) (LLVM.Value (Vector n a))


interleave ::
   (LLVM.IsPrimitive a, TypeNum.Pos n) =>
   Stereo.T (LLVM.Value (Vector n a)) ->
   LLVM.CodeGenFunction r (Value n a)
interleave x =
   uncurry (liftM2 Value) .
   mapPair (Vector.assemble, Vector.assemble) .
   splitAt (Vector.sizeInTuple x) .
   concatMap (\s -> [Stereo.left s, Stereo.right s]) =<<
   Vector.extractAll x

deinterleave ::
   (LLVM.IsPrimitive a, TypeNum.Pos n) =>
   Value n a ->
   LLVM.CodeGenFunction r (Stereo.T (LLVM.Value (Vector n a)))
deinterleave (Value v0 v1) =
   Vector.assemble .
   (let aux (l:r:xs) = Stereo.cons l r : aux xs
        aux [] = []
        aux _ = error "odd number of stereo elements"
    in  aux) =<<
   liftM2 (++)
      (Vector.extractAll v0)
      (Vector.extractAll v1)


instance
   (TypeNum.Pos n, LLVM.IsPrimitive a, St.Storable a) =>
      St.Storable (T n a) where
   sizeOf ~(Cons v0 v1) = St.sizeOf v0 + St.sizeOf v1
   alignment ~(Cons v _) = St.alignment v
   peek ptr =
      let p = castPtr ptr
      in  liftM2 Cons
             (St.peekElemOff p 0)
             (St.peekElemOff p 1)
   poke ptr (Cons v0 v1) =
      let p = castPtr ptr
      in  St.pokeElemOff p 0 v0 >>
          St.pokeElemOff p 1 v1

instance (TypeNum.Pos n, LLVM.IsPrimitive a) => Class.Zero (Value n a) where
   zeroTuple = Value (LLVM.value LLVM.zero) (LLVM.value LLVM.zero)

instance (TypeNum.Pos n, LLVM.IsPrimitive a) => Undefined (Value n a) where
   undefTuple = Value (LLVM.value LLVM.undef) (LLVM.value LLVM.undef)

{-
Can only be implemented by ifThenElse
since the atomic 'select' command wants a bool vector.

instance (TypeNum.Pos n, LLVM.IsPrimitive a, Phi a) => C.Select (Value n a) where
   select b (Value x0 x1) (Value y0 y1) =
      liftM2 Value
         (C.select b x0 y0)
         (C.select b x1 y1)

instance LLVM.CmpRet a b => LLVM.CmpRet (Stereo.T a) (Stereo.T b) where
-}

instance (TypeNum.Pos n, LLVM.IsPrimitive a, LLVM.IsConst a) =>
      MakeValueTuple (T n a) (Value n a) where
   valueTupleOf (Cons v0 v1) =
      Value
         (LLVM.valueOf v0)
         (LLVM.valueOf v1)

instance (TypeNum.Pos n, LLVM.IsPrimitive a) => Phi (Value n a) where
   phis bb = mapV (phis bb)
   addPhis bb = zipV (\_ _ -> ()) (addPhis bb)

{-
instance Vector.ShuffleMatch n (Value n a) where
   shuffleMatch = Vector.shuffleMatchAccess
-}

{-
We cannot make an instance of Vector.Access,
since we defined the functional dependencies in a way,
that vector size and element type uniquely defines the vector type.

instance Vector.Access n (Stereo.T a) (Value n a) where

insert ::
   LLVM.Value Word32 ->
   LLVM.Value a ->
   Value n a ->
   LLVM.CodeGenFunction r (Value n a)
insert k a (Value v0 v1) = do
   k20 <- A.add k k
   k21 <- A.inc k20
   select (k20<Vector.size v0) ...
   Vector.insert k20 (Stereo.left a) v >>=
      Vector.insert k21 (Stereo.right a)

extract ::
   LLVM.Value Word32 ->
   Value n a ->
   LLVM.CodeGenFunction r (LLVM.Value a)
extract k (Value v0 v1) = do
   k20 <- A.add k k
   k21 <- A.inc k20
   liftM2 Value
      (Vector.extract k20 v)
      (Vector.extract k21 v)


With this instance,
both 'interleave' and 'deinterleave' could then be written this way:

   Vector.assemble =<<
   mapM (flip Vector.extract x) (take (Vector.size x) [0..])
-}

memory ::
   (TypeNum.Pos n, LLVM.IsPrimitive a, LLVM.IsPrimitive am,
    Memory.FirstClass a am) =>
   Memory.Record r (Struct (Vector n am, (Vector n am, ()))) (Value n a)
memory =
   liftA2 Value
      (Memory.element (\(Value v _) -> v) TypeNum.d0)
      (Memory.element (\(Value _ v) -> v) TypeNum.d1)

instance
      (TypeNum.Pos n,
       Memory.FirstClass a am,
       LLVM.IsPrimitive a, IsSized a as,
       TypeNum.Mul n as vs, TypeNum.Pos vs,
       LLVM.IsPrimitive am, IsSized am amsize,
       TypeNum.Mul n amsize vmsize, TypeNum.Pos vmsize) =>
      Memory.C (Value n a) (Struct (Vector n am, (Vector n am, ()))) where
   load = Memory.loadRecord memory
   store = Memory.storeRecord memory
   decompose = Memory.decomposeRecord memory
   compose = Memory.composeRecord memory

{- |
This instance allows to run @arrange@ on interleaved stereo vectors.
-}
instance
   (TypeNum.Pos n, LLVM.IsPrimitive a, LLVM.IsArithmetic a) =>
      A.Additive (Value n a) where
   zero = Value A.zero A.zero
   add = zipV Value A.add
   sub = zipV Value A.sub
   neg = mapV A.neg


mapV :: (Monad m) =>
   (LLVM.Value (Vector n a) -> m (LLVM.Value (Vector n a))) ->
   Value n a -> m (Value n a)
mapV f (Value x0 x1) =
   liftM2 Value (f x0) (f x1)

zipV :: (Monad m) =>
   (c -> c -> d) ->
   (LLVM.Value (Vector n a) ->
    LLVM.Value (Vector n b) ->
    m c) ->
   Value n a ->
   Value n b ->
   m d
zipV g f (Value x0 x1) (Value y0 y1) =
   liftM2 g (f x0 y0) (f x1 y1)