{-# LANGUAGE TypeFamilies #-}
module LLVM.Extra.Multi.Vector where

import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Vector as Vector
import qualified LLVM.Extra.Class as Class

import qualified LLVM.Core as LLVM
import LLVM.Core
   (Value, ConstValue, valueOf, value,
    IsPrimitive,
    CodeGenFunction, )

import qualified Types.Data.Num as TypeNum

import qualified Data.List as List
import Data.Word (Word32, )

import Control.Monad (liftM2, liftM3, foldM, )

import Prelude hiding (zip, zip3, unzip, unzip3, )


newtype T n a = Cons (Vector n a)

instance (TypeNum.PositiveT n, C a) => Class.Undefined (T n a) where
   undefTuple = undef


size :: TypeNum.PositiveT n => T n a -> Int
size =
   let sz :: TypeNum.PositiveT n => n -> T n a -> Int
       sz n _ = TypeNum.fromIntegerT n
   in  sz undefined


zip :: T n a -> T n b -> T n (a,b)
zip (Cons a) (Cons b) = Cons (a,b)

zip3 :: T n a -> T n b -> T n c -> T n (a,b,c)
zip3 (Cons a) (Cons b) (Cons c) = Cons (a,b,c)

unzip :: T n (a,b) -> (T n a, T n b)
unzip (Cons (a,b)) = (Cons a, Cons b)

unzip3 :: T n (a,b,c) -> (T n a, T n b, T n c)
unzip3 (Cons (a,b,c)) = (Cons a, Cons b, Cons c)


class C a where
   type Vector n a :: *
   undef :: (TypeNum.PositiveT n) => T n a
   shuffleMatch ::
      (TypeNum.PositiveT n) =>
      ConstValue (LLVM.Vector n Word32) -> T n a -> CodeGenFunction r (T n a)
   extract ::
      (TypeNum.PositiveT n) =>
      LLVM.Value Word32 -> T n a -> CodeGenFunction r (MultiValue.T a)
   insert ::
      (TypeNum.PositiveT n) =>
      LLVM.Value Word32 -> MultiValue.T a ->
      T n a -> CodeGenFunction r (T n a)

instance C Float where
   type Vector n Float = LLVM.Value (LLVM.Vector n Float)
   undef = undefPrimitive
   shuffleMatch = shuffleMatchPrimitive
   extract = extractPrimitive
   insert = insertPrimitive

instance C Double where
   type Vector n Double = LLVM.Value (LLVM.Vector n Double)
   undef = undefPrimitive
   shuffleMatch = shuffleMatchPrimitive
   extract = extractPrimitive
   insert = insertPrimitive

undefPrimitive ::
   (TypeNum.PositiveT n, IsPrimitive a,
    Vector n a ~ Value (LLVM.Vector n a)) =>
   T n a
undefPrimitive = Cons $ LLVM.value LLVM.undef

shuffleMatchPrimitive ::
   (TypeNum.PositiveT n, IsPrimitive a,
    Vector n a ~ Value (LLVM.Vector n a),
    Class.ValueTuple a ~ Value a) =>
   ConstValue (LLVM.Vector n Word32) -> T n a -> CodeGenFunction r (T n a)
shuffleMatchPrimitive k (Cons v) =
   fmap Cons $ LLVM.shufflevector v (value LLVM.undef) k

extractPrimitive ::
   (TypeNum.PositiveT n, IsPrimitive a,
    Vector n a ~ Value (LLVM.Vector n a),
    Class.ValueTuple a ~ Value a) =>
   Value Word32 -> T n a -> CodeGenFunction r (MultiValue.T a)
extractPrimitive k (Cons v) =
   fmap MultiValue.Cons $ LLVM.extractelement v k

insertPrimitive ::
   (TypeNum.PositiveT n, IsPrimitive a,
    Vector n a ~ Value (LLVM.Vector n a),
    Class.ValueTuple a ~ Value a) =>
   Value Word32 ->
   MultiValue.T a -> T n a -> CodeGenFunction r (T n a)
insertPrimitive k (MultiValue.Cons a) (Cons v) =
   fmap Cons $ LLVM.insertelement v a k


instance (C a, C b) => C (a,b) where
   type Vector n (a,b) = (Vector n a, Vector n b)

   undef = zip undef undef

   shuffleMatch is v =
      case unzip v of
         (v0,v1) ->
            liftM2 zip
               (shuffleMatch is v0)
               (shuffleMatch is v1)

   extract k v =
      case unzip v of
         (v0,v1) ->
            liftM2 MultiValue.zip
               (extract k v0)
               (extract k v1)

   insert k a v =
      case (MultiValue.unzip a, unzip v) of
         ((a0,a1), (v0,v1)) ->
            liftM2 zip
               (insert k a0 v0)
               (insert k a1 v1)


instance (C a, C b, C c) => C (a,b,c) where
   type Vector n (a,b,c) = (Vector n a, Vector n b, Vector n c)

   undef = zip3 undef undef undef

   shuffleMatch is v =
      case unzip3 v of
         (v0,v1,v2) ->
            liftM3 zip3
               (shuffleMatch is v0)
               (shuffleMatch is v1)
               (shuffleMatch is v2)

   extract k v =
      case unzip3 v of
         (v0,v1,v2) ->
            liftM3 MultiValue.zip3
               (extract k v0)
               (extract k v1)
               (extract k v2)

   insert k a v =
      case (MultiValue.unzip3 a, unzip3 v) of
         ((a0,a1,a2), (v0,v1,v2)) ->
            liftM3 zip3
               (insert k a0 v0)
               (insert k a1 v1)
               (insert k a2 v2)


assemble ::
   (TypeNum.PositiveT n, C a) =>
   [MultiValue.T a] -> CodeGenFunction r (T n a)
assemble =
   foldM (\v (k,x) -> insert (valueOf k) x v) undef .
   List.zip [0..]

dissect ::
   (TypeNum.PositiveT n, C a) =>
   T n a -> LLVM.CodeGenFunction r [MultiValue.T a]
dissect x =
   mapM
      (flip extract x . LLVM.valueOf)
      (take (size x) [0..])


-- * function based on classes from Vector module

shuffleMatchGen ::
   (Vector n a ~ v, Vector.Simple v, n ~ Vector.Size v) =>
   ConstValue (LLVM.Vector n Word32) ->
   T n a -> CodeGenFunction r (T n a)
shuffleMatchGen is (Cons v) =
   fmap Cons $ Vector.shuffleMatch is v

extractGen ::
   (Vector n a ~ v, Vector.Simple v, Class.ValueTuple a ~ Vector.Element v) =>
   LLVM.Value Word32 -> T n a -> CodeGenFunction r (MultiValue.T a)
extractGen n (Cons v) =
   fmap MultiValue.Cons $ Vector.extract n v

insertGen ::
   (Vector n a ~ v, Vector.C v, Class.ValueTuple a ~ Vector.Element v) =>
   LLVM.Value Word32 -> MultiValue.T a ->
   T n a -> CodeGenFunction r (T n a)
insertGen n (MultiValue.Cons a) (Cons v) =
   fmap Cons $ Vector.insert n a v