Vector.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module LLVM.DSL.Expression.Vector where

import qualified LLVM.DSL.Expression as Expr
import LLVM.DSL.Expression (Exp)

import qualified LLVM.Extra.Nice.Value.Vector as NiceValueVec
import qualified LLVM.Extra.Nice.Value as NiceValue
import qualified LLVM.Extra.Nice.Vector.Instance as NiceVectorInst
import qualified LLVM.Extra.Nice.Vector as NiceVector
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Core as LLVM

import qualified Data.Tuple.HT as Tuple

import Prelude hiding (replicate, take, zip, fst, snd, min, max)


cons ::
   (LLVM.Positive n, NiceVector.C a) =>
   LLVM.Vector n a -> Exp (LLVM.Vector n a)
cons = Expr.lift0 . NiceValueVec.cons

fst ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n a)
fst = Expr.lift1 NiceValueVec.fst

snd ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n b)
snd = Expr.lift1 NiceValueVec.snd

swap ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n (b,a))
swap = Expr.lift1 NiceValueVec.swap

mapFst ::
   (Exp (LLVM.Vector n a0) -> Exp (LLVM.Vector n a1)) ->
   Exp (LLVM.Vector n (a0,b)) -> Exp (LLVM.Vector n (a1,b))
mapFst f =
   Expr.liftReprM
      (\(a0,b) -> do
         NiceValue.Cons a1 <- Expr.unliftM1 f $ NiceValue.Cons a0
         return (a1,b))

mapSnd ::
   (Exp (LLVM.Vector n b0) -> Exp (LLVM.Vector n b1)) ->
   Exp (LLVM.Vector n (a,b0)) -> Exp (LLVM.Vector n (a,b1))
mapSnd f =
   Expr.liftReprM
      (\(a,b0) -> do
         NiceValue.Cons b1 <- Expr.unliftM1 f $ NiceValue.Cons b0
         return (a,b1))


fst3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n a)
fst3 = Expr.lift1 NiceValueVec.fst3

snd3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n b)
snd3 = Expr.lift1 NiceValueVec.snd3

thd3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n c)
thd3 = Expr.lift1 NiceValueVec.thd3


zip ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n b) ->
   Exp (LLVM.Vector n (a,b))
zip = Expr.lift2 NiceValueVec.zip

zip3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n b) -> Exp (LLVM.Vector n c) ->
   Exp (LLVM.Vector n (a,b,c))
zip3 = Expr.lift3 NiceValueVec.zip3


replicate ::
   (LLVM.Positive n, NiceVector.C a) =>
   Exp a -> Exp (LLVM.Vector n a)
replicate = Expr.liftM NiceValueVec.replicate

iterate ::
   (LLVM.Positive n, NiceVector.C a) =>
   (Exp a -> Exp a) -> Exp a -> Exp (LLVM.Vector n a)
iterate f = Expr.liftM (NiceValueVec.iterate (Expr.unliftM1 f))

take ::
   (LLVM.Positive n, LLVM.Positive m, NiceVector.Select a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector m a)
take = Expr.liftM NiceValueVec.take

takeRev ::
   (LLVM.Positive n, LLVM.Positive m, NiceVector.Select a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector m a)
takeRev = Expr.liftM NiceValueVec.takeRev


cumulate ::
   (LLVM.Positive n, NiceVector.Additive a) =>
   Exp a -> Exp (LLVM.Vector n a) -> (Exp a, Exp (LLVM.Vector n a))
cumulate a0 v0 =
   Expr.unzip $
   Expr.liftM2
      (\a v ->
         fmap (uncurry NiceValue.zip .
               Tuple.mapSnd NiceVectorInst.toNiceValue) $
         NiceVector.cumulate a $ NiceVectorInst.fromNiceValue v)
      a0 v0


cmp ::
   (LLVM.Positive n, NiceVector.Comparison a) =>
   LLVM.CmpPredicate ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n Bool)
cmp ord = Expr.liftM2 (NiceValueVec.cmp ord)

select ::
   (LLVM.Positive n, NiceVector.Select a) =>
   Exp (LLVM.Vector n Bool) ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
select = Expr.liftM3 NiceValueVec.select


min, max ::
   (LLVM.Positive n, NiceVector.Real a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
min = Expr.liftM2 A.min
max = Expr.liftM2 A.max

limit ::
   (LLVM.Positive n, NiceVector.Real a) =>
   (Exp (LLVM.Vector n a), Exp (LLVM.Vector n a)) ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
limit (l,u) = max l . min u


fromIntegral ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp i -> Exp a
fromIntegral = Expr.liftM NiceValueVec.fromIntegral

truncateToInt ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp a -> Exp i
truncateToInt = Expr.liftM NiceValueVec.truncateToInt

splitFractionToInt ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp a -> (Exp i, Exp a)
splitFractionToInt = Expr.unzip . Expr.liftM NiceValueVec.splitFractionToInt