module LLVM.DSL.Example.Median where

import qualified LLVM.DSL.Expression as Expr
import LLVM.DSL.Expression (Exp, (==*), (<=*), (&&*), (||*))

import qualified LLVM.Extra.Multi.Vector as MVec
import qualified LLVM.Extra.Multi.Value as MV

import qualified LLVM.Core as LLVM

import Data.Tuple.HT (uncurry3)
import Data.Word (Word8)



median3IfThen :: (MV.Comparison a) => Exp a -> Exp a -> Exp a -> Exp a
median3IfThen :: forall a. Comparison a => Exp a -> Exp a -> Exp a -> Exp a
median3IfThen Exp a
a Exp a
b Exp a
c =
   Exp Bool -> Exp a -> Exp a -> Exp a
forall a. C a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.ifThenElse (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
b)
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. C a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.ifThenElse (Exp a
bExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c)
         Exp a
b
         (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. C a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.ifThenElse (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) Exp a
c Exp a
a))
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. C a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.ifThenElse (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c)
         Exp a
a
         (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. C a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.ifThenElse (Exp a
bExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) Exp a
c Exp a
b))

median3Select ::
   (MV.Comparison a, MV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3Select :: forall a.
(Comparison a, Select a) =>
Exp a -> Exp a -> Exp a -> Exp a
median3Select Exp a
a Exp a
b Exp a
c =
   Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
b)
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
bExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c)
         Exp a
b
         (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) Exp a
c Exp a
a))
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c)
         Exp a
a
         (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
bExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) Exp a
c Exp a
b))

median3SelectShared ::
   (MV.Comparison a, MV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3SelectShared :: forall a.
(Comparison a, Select a) =>
Exp a -> Exp a -> Exp a -> Exp a
median3SelectShared Exp a
a Exp a
b Exp a
c =
   Exp Bool -> (Exp Bool -> Exp a) -> Exp a
forall a b. Exp a -> (Exp a -> Exp b) -> Exp b
Expr.with (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
b) ((Exp Bool -> Exp a) -> Exp a) -> (Exp Bool -> Exp a) -> Exp a
forall a b. (a -> b) -> a -> b
$ \Exp Bool
a_le_b ->
   Exp Bool -> (Exp Bool -> Exp a) -> Exp a
forall a b. Exp a -> (Exp a -> Exp b) -> Exp b
Expr.with (Exp a
bExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) ((Exp Bool -> Exp a) -> Exp a) -> (Exp Bool -> Exp a) -> Exp a
forall a b. (a -> b) -> a -> b
$ \Exp Bool
b_le_c ->
   Exp Bool -> (Exp Bool -> Exp a) -> Exp a
forall a b. Exp a -> (Exp a -> Exp b) -> Exp b
Expr.with (Exp a
aExp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=*Exp a
c) ((Exp Bool -> Exp a) -> Exp a) -> (Exp Bool -> Exp a) -> Exp a
forall a b. (a -> b) -> a -> b
$ \Exp Bool
a_le_c ->
   Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select Exp Bool
a_le_b
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select Exp Bool
b_le_c Exp a
b (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select Exp Bool
a_le_c Exp a
c Exp a
a))
      (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select Exp Bool
a_le_c Exp a
a (Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select Exp Bool
b_le_c Exp a
c Exp a
b))

median3MinMax ::
   (MV.Comparison a, MV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3MinMax :: forall a.
(Comparison a, Select a) =>
Exp a -> Exp a -> Exp a -> Exp a
median3MinMax Exp a
a Exp a
b Exp a
c =
   let minab :: Exp a
minab = Exp a -> Exp a -> Exp a
forall a. Real a => Exp a -> Exp a -> Exp a
Expr.min Exp a
a Exp a
b in
   let maxab :: Exp a
maxab = Exp a -> Exp a -> Exp a
forall a. Real a => Exp a -> Exp a -> Exp a
Expr.max Exp a
a Exp a
b in
   Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
maxab Exp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=* Exp a
c) Exp a
maxab (Exp a -> Exp a) -> Exp a -> Exp a
forall a b. (a -> b) -> a -> b
$ Exp Bool -> Exp a -> Exp a -> Exp a
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (Exp a
minab Exp a -> Exp a -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
<=* Exp a
c) Exp a
c Exp a
minab

median3MinMaxVector ::
   (LLVM.Positive n, MVec.C a) =>
   (MV.Comparison a, MV.Select a) =>
   MVec.T n a -> MVec.T n a -> MVec.T n a ->
   LLVM.CodeGenFunction r (MVec.T n a)
median3MinMaxVector :: forall n a r.
(Positive n, C a, Comparison a, Select a) =>
T n a -> T n a -> T n a -> CodeGenFunction r (T n a)
median3MinMaxVector T n a
a T n a
b T n a
c =
   (T (a, a, a) -> CodeGenFunction r (T a))
-> T n (a, a, a) -> CodeGenFunction r (T n a)
forall n a b r.
(Positive n, C a, C b) =>
(T a -> CodeGenFunction r (T b))
-> T n a -> CodeGenFunction r (T n b)
MVec.map ((T a -> T a -> T a -> CodeGenFunction r (T a))
-> (T a, T a, T a) -> CodeGenFunction r (T a)
forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 ((Exp a -> Exp a -> Exp a -> Exp a)
-> T a -> T a -> T a -> CodeGenFunction r (T a)
forall ae am be bm ce cm de dm r.
(Aggregate ae am, Aggregate be bm, Aggregate ce cm,
 Aggregate de dm) =>
(ae -> be -> ce -> de) -> am -> bm -> cm -> CodeGenFunction r dm
Expr.unliftM3 Exp a -> Exp a -> Exp a -> Exp a
forall a.
(Comparison a, Select a) =>
Exp a -> Exp a -> Exp a -> Exp a
median3MinMax) ((T a, T a, T a) -> CodeGenFunction r (T a))
-> (T (a, a, a) -> (T a, T a, T a))
-> T (a, a, a)
-> CodeGenFunction r (T a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (a, a, a) -> (T a, T a, T a)
forall a b c. T (a, b, c) -> (T a, T b, T c)
MV.unzip3) (T n (a, a, a) -> CodeGenFunction r (T n a))
-> T n (a, a, a) -> CodeGenFunction r (T n a)
forall a b. (a -> b) -> a -> b
$
   T n a -> T n a -> T n a -> T n (a, a, a)
forall n a b c. T n a -> T n b -> T n c -> T n (a, b, c)
MVec.zip3 T n a
a T n a
b T n a
c


type MV = MV.T

median3Case ::
   (MV.Comparison a, MV.Select a) =>
   MV a -> MV a -> MV a -> LLVM.CodeGenFunction r (MV a)
median3Case :: forall a r.
(Comparison a, Select a) =>
MV a -> MV a -> MV a -> CodeGenFunction r (MV a)
median3Case MV a
a MV a
b MV a
c = do
   T Bool
a_le_b <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
a MV a
b
   T Bool
a_le_c <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
a MV a
c
   T Bool
b_le_c <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
b MV a
c
   let mask :: Integer -> T Word8
mask = Integer -> T Word8
forall a. IntegerConstant a => Integer -> T a
MV.fromInteger'
   T Word8
a_le_b_mask <- T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r.
T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
MV.select T Bool
a_le_b (Integer -> T Word8
mask Integer
1) (Integer -> T Word8
mask Integer
0)
   T Word8
a_le_c_mask <- T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r.
T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
MV.select T Bool
a_le_c (Integer -> T Word8
mask Integer
2) (Integer -> T Word8
mask Integer
0)
   T Word8
b_le_c_mask <- T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r.
T Bool -> T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
MV.select T Bool
b_le_c (Integer -> T Word8
mask Integer
4) (Integer -> T Word8
mask Integer
0)
   T Word8
maskMV <- T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
forall a r. Logic a => T a -> T a -> CodeGenFunction r (T a)
forall r. T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
MV.or T Word8
a_le_b_mask (T Word8 -> CodeGenFunction r (T Word8))
-> CodeGenFunction r (T Word8) -> CodeGenFunction r (T Word8)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
forall a r. Logic a => T a -> T a -> CodeGenFunction r (T a)
forall r. T Word8 -> T Word8 -> CodeGenFunction r (T Word8)
MV.or T Word8
a_le_c_mask T Word8
b_le_c_mask
   let maskE :: Exp Word8
maskE = T Word8 -> Exp Word8
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T Word8
maskMV :: MV Word8)
   T Bool
selectB <- Exp Bool -> forall r. CodeGenFunction r (T Bool)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
Expr.unExp (Exp Word8
maskE Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* Exp Word8
0 Exp Bool -> Exp Bool -> Exp Bool
||* Exp Word8
maskE Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* Exp Word8
7)
   T Bool
selectA <- Exp Bool -> forall r. CodeGenFunction r (T Bool)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
Expr.unExp (Exp Word8
maskE Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* Exp Word8
1 Exp Bool -> Exp Bool -> Exp Bool
||* Exp Word8
maskE Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* Exp Word8
6)
   T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r. T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
MV.select T Bool
selectA MV a
a (MV a -> CodeGenFunction r (MV a))
-> CodeGenFunction r (MV a) -> CodeGenFunction r (MV a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r. T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
MV.select T Bool
selectB MV a
b MV a
c

median3CaseVec ::
   (MV.Comparison a, MV.Select a) =>
   MV a -> MV a -> MV a -> LLVM.CodeGenFunction r (MV a)
median3CaseVec :: forall a r.
(Comparison a, Select a) =>
MV a -> MV a -> MV a -> CodeGenFunction r (MV a)
median3CaseVec MV a
a MV a
b MV a
c = do
   T Bool
a_le_b <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
a MV a
b
   T Bool
a_le_c <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
a MV a
c
   T Bool
b_le_c <- CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> MV a -> MV a -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MV.cmp CmpPredicate
LLVM.CmpLE MV a
b MV a
c
   let check :: Exp Word8 -> Exp Word8 -> Exp Word8 -> Exp Bool
check Exp Word8
ab Exp Word8
ac Exp Word8
bc =
         Exp Bool -> Exp Word8 -> Exp Word8 -> Exp Word8
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (T Bool -> Exp Bool
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 T Bool
a_le_b) Exp Word8
1 Exp Word8
0 Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* (Exp Word8
ab :: Exp Word8)
         Exp Bool -> Exp Bool -> Exp Bool
&&*
         Exp Bool -> Exp Word8 -> Exp Word8 -> Exp Word8
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (T Bool -> Exp Bool
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 T Bool
a_le_c) Exp Word8
1 Exp Word8
0 Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* (Exp Word8
ac :: Exp Word8)
         Exp Bool -> Exp Bool -> Exp Bool
&&*
         Exp Bool -> Exp Word8 -> Exp Word8 -> Exp Word8
forall a. Select a => Exp Bool -> Exp a -> Exp a -> Exp a
Expr.select (T Bool -> Exp Bool
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 T Bool
b_le_c) Exp Word8
1 Exp Word8
0 Exp Word8 -> Exp Word8 -> Exp Bool
forall a. Comparison a => Exp a -> Exp a -> Exp Bool
==* (Exp Word8
bc :: Exp Word8)
   T Bool
selectB <- Exp Bool -> forall r. CodeGenFunction r (T Bool)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
Expr.unExp (Exp Word8 -> Exp Word8 -> Exp Word8 -> Exp Bool
check Exp Word8
0 Exp Word8
0 Exp Word8
0 Exp Bool -> Exp Bool -> Exp Bool
||* Exp Word8 -> Exp Word8 -> Exp Word8 -> Exp Bool
check Exp Word8
1 Exp Word8
1 Exp Word8
1)
   T Bool
selectA <- Exp Bool -> forall r. CodeGenFunction r (T Bool)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
Expr.unExp (Exp Word8 -> Exp Word8 -> Exp Word8 -> Exp Bool
check Exp Word8
1 Exp Word8
0 Exp Word8
0 Exp Bool -> Exp Bool -> Exp Bool
||* Exp Word8 -> Exp Word8 -> Exp Word8 -> Exp Bool
check Exp Word8
0 Exp Word8
1 Exp Word8
1)
   T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r. T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
MV.select T Bool
selectA MV a
a (MV a -> CodeGenFunction r (MV a))
-> CodeGenFunction r (MV a) -> CodeGenFunction r (MV a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
forall a r.
Select a =>
T Bool -> T a -> T a -> CodeGenFunction r (T a)
forall r. T Bool -> MV a -> MV a -> CodeGenFunction r (MV a)
MV.select T Bool
selectB MV a
b MV a
c