module Futhark.Analysis.PrimExp.TableTests (tests) where

import Control.Monad.State.Strict
import Data.Map.Strict qualified as M
import Futhark.Analysis.PrimExp
import Futhark.Analysis.PrimExp.Table
import Futhark.IR.GPU
import Futhark.IR.GPUTests ()
import Futhark.IR.MC
import Futhark.IR.MCTests ()
import Test.Tasty
import Test.Tasty.HUnit

tests :: TestTree
tests :: TestTree
tests = TestName -> [TestTree] -> TestTree
testGroup TestName
"AnalyzePrim" [TestTree
stmToPrimExpsTests]

stmToPrimExpsTests :: TestTree
stmToPrimExpsTests :: TestTree
stmToPrimExpsTests =
  TestName -> [TestTree] -> TestTree
testGroup
    TestName
"stmToPrimExps"
    [TestTree
stmToPrimExpsTestsGPU, TestTree
stmToPrimExpsTestsMC]

stmToPrimExpsTestsGPU :: TestTree
stmToPrimExpsTestsGPU :: TestTree
stmToPrimExpsTestsGPU =
  TestName -> [TestTree] -> TestTree
testGroup
    TestName
"GPU"
    ([TestTree] -> TestTree) -> [TestTree] -> TestTree
forall a b. (a -> b) -> a -> b
$ do
      let scope :: Map VName (NameInfo GPU)
scope =
            [(VName, NameInfo GPU)] -> Map VName (NameInfo GPU)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
              [ (VName
"n_5142", FParamInfo GPU -> NameInfo GPU
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo GPU
"i64"),
                (VName
"m_5143", FParamInfo GPU -> NameInfo GPU
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo GPU
"i64"),
                (VName
"xss_5144", FParamInfo GPU -> NameInfo GPU
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo GPU
"[n_5142][m_5143]i64"),
                (VName
"segmap_group_size_5201", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"i64"),
                (VName
"segmap_usable_groups_5202", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"i64"),
                (VName
"defunc_0_map_res_5203", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"[n_5142]i64"),
                (VName
"defunc_0_f_res_5207", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"i64"),
                (VName
"i_5208", IntType -> NameInfo GPU
forall rep. IntType -> NameInfo rep
IndexName IntType
Int64),
                (VName
"acc_5209", FParamInfo GPU -> NameInfo GPU
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo GPU
"i64"),
                (VName
"b_5210", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"i64"),
                (VName
"defunc_0_f_res_5211", LetDec GPU -> NameInfo GPU
forall rep. LetDec rep -> NameInfo rep
LetName LetDec GPU
"i64")
              ]
      [ TestName -> Assertion -> TestTree
testCase TestName
"BinOp" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm GPU
stm = Stm GPU
"let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ ( VName
"defunc_0_f_res_5211",
                      PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                        ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"b_5210" (IntType -> PrimType
IntType IntType
Int64))
                        )
                    )
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Index" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm GPU
stm = Stm GPU
"let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: Map VName (Maybe a)
expected = [(VName, Maybe a)] -> Map VName (Maybe a)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName
"b_5210", Maybe a
forall a. Maybe a
Nothing)]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
forall {a}. Map VName (Maybe a)
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Loop" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm GPU
stm = Stm GPU
"let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { {defunc_0_f_res_5211} }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_f_res_5207", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"i_5208", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"i_5208" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"acc_5209", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64)))
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Loop body" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm GPU
stm = Stm GPU
"let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { let {b_5210 : i64} = xss_5144[gtid_5204, i_5208] let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210) in {defunc_0_f_res_5211} }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_f_res_5207", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"i_5208", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"i_5208" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"acc_5209", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"b_5210", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    ( VName
"defunc_0_f_res_5211",
                      PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                        ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"b_5210" (IntType -> PrimType
IntType IntType
Int64))
                        )
                    )
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"SegMap" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$
          do
            let stm :: Stm GPU
stm =
                  Stm GPU
"let {defunc_0_map_res_5125 : [n_5142]i64} =\
                  \  segmap(thread; ; grid=segmap_usable_groups_5124; blocksize=segmap_group_size_5123)\
                  \  (gtid_5126 < n_5142) (~phys_tid_5127) : {i64} {\
                  \  return {returns lifted_lambda_res_5129} \
                  \}"
            let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
            let expected :: PrimExpTable
expected =
                  [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                    [ (VName
"defunc_0_map_res_5125", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                      (VName
"gtid_5126", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"gtid_5126" (IntType -> PrimType
IntType IntType
Int64)))
                    ]
            PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"SegMap body" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$
          do
            let stm :: Stm GPU
                stm :: Stm GPU
stm =
                  Stm GPU
"let {defunc_0_map_res_5125 : [n_5142]i64} =\
                  \  segmap(thread; ; grid=segmap_usable_groups_5124; blocksize=segmap_group_size_5123)\
                  \  (gtid_5126 < n_5142) (~phys_tid_5127) : {i64} {\
                  \    let {eta_p_5128 : i64} =\
                  \      xs_5093[gtid_5126]\
                  \    let {lifted_lambda_res_5129 : i64} =\
                  \      add64(2i64, eta_p_5128)\
                  \    return {returns lifted_lambda_res_5129}\
                  \  }"
            let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo GPU) -> Stm GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo GPU)
scope Stm GPU
stm) PrimExpTable
forall a. Monoid a => a
mempty
            let expected :: PrimExpTable
expected =
                  [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                    [ (VName
"defunc_0_map_res_5125", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                      (VName
"gtid_5126", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"gtid_5126" (IntType -> PrimType
IntType IntType
Int64))),
                      (VName
"eta_p_5128", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                      ( VName
"lifted_lambda_res_5129",
                        PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                          ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                              (PrimValue -> PrimExp VName
forall v. PrimValue -> PrimExp v
ValueExp (IntValue -> PrimValue
IntValue (Int64 -> IntValue
Int64Value Int64
2)))
                              (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"eta_p_5128" (IntType -> PrimType
IntType IntType
Int64))
                          )
                      )
                    ]
            PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected
        ]

stmToPrimExpsTestsMC :: TestTree
stmToPrimExpsTestsMC :: TestTree
stmToPrimExpsTestsMC =
  TestName -> [TestTree] -> TestTree
testGroup
    TestName
"MC"
    ([TestTree] -> TestTree) -> [TestTree] -> TestTree
forall a b. (a -> b) -> a -> b
$ do
      let scope :: Map VName (NameInfo MC)
scope =
            [(VName, NameInfo MC)] -> Map VName (NameInfo MC)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
              [ (VName
"n_5142", FParamInfo MC -> NameInfo MC
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo MC
"i64"),
                (VName
"m_5143", FParamInfo MC -> NameInfo MC
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo MC
"i64"),
                (VName
"xss_5144", FParamInfo MC -> NameInfo MC
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo MC
"[n_5142][5143]i64"),
                (VName
"segmap_group_size_5201", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"i64"),
                (VName
"segmap_usable_groups_5202", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"i64"),
                (VName
"defunc_0_map_res_5203", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"[n_5142]i64"),
                (VName
"defunc_0_f_res_5207", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"i64"),
                (VName
"i_5208", IntType -> NameInfo MC
forall rep. IntType -> NameInfo rep
IndexName IntType
Int64),
                (VName
"acc_5209", FParamInfo MC -> NameInfo MC
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo MC
"i64"),
                (VName
"b_5210", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"i64"),
                (VName
"defunc_0_f_res_5211", LetDec MC -> NameInfo MC
forall rep. LetDec rep -> NameInfo rep
LetName LetDec MC
"i64")
              ]
      [ TestName -> Assertion -> TestTree
testCase TestName
"BinOp" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
stm = Stm MC
"let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ ( VName
"defunc_0_f_res_5211",
                      PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                        ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"b_5210" (IntType -> PrimType
IntType IntType
Int64))
                        )
                    )
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Index" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
stm = Stm MC
"let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: Map VName (Maybe a)
expected = [(VName, Maybe a)] -> Map VName (Maybe a)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName
"b_5210", Maybe a
forall a. Maybe a
Nothing)]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
forall {a}. Map VName (Maybe a)
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Loop" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
stm = Stm MC
"let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { {defunc_0_f_res_5211} }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_f_res_5207", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"i_5208", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"i_5208" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"acc_5209", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64)))
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"Loop body" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
stm =
                Stm MC
"\
                \let {defunc_0_f_res_5207 : i64} =\
                \  loop {acc_5209 : i64} = {0i64}\
                \  for i_5208:i64 < m_5143 do {\
                \    let {b_5210 : i64} =\
                \      xss_5144[gtid_5204, i_5208]\
                \    let {defunc_0_f_res_5211 : i64} =\
                \      add64(acc_5209, b_5210)\
                \    in {defunc_0_f_res_5211}\
                \  }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_f_res_5207", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"i_5208", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"i_5208" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"acc_5209", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"b_5210", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    ( VName
"defunc_0_f_res_5211",
                      PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                        ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"acc_5209" (IntType -> PrimType
IntType IntType
Int64))
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"b_5210" (IntType -> PrimType
IntType IntType
Int64))
                        )
                    )
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"SegMap" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
stm =
                Stm MC
"let {defunc_0_map_res_5125 : [n_5142]i64} =\
                \  segmap()\
                \  (gtid_5126 < n_5142) (~flat_tid_5112) : {i64} {\
                \    return {returns lifted_lambda_res_5129}\
                \  }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_map_res_5125", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"gtid_5126", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"gtid_5126" (IntType -> PrimType
IntType IntType
Int64)))
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected,
        TestName -> Assertion -> TestTree
testCase TestName
"SegMap body" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
          let stm :: Stm MC
              stm :: Stm MC
stm =
                Stm MC
"let {defunc_0_map_res_5125 : [n_5142]i64} =\
                \  segmap()\
                \  (gtid_5126 < n_5142) (~flat_tid_5112) : {i64} {\
                \    let {eta_p_5128 : i64} =\
                \      xs_5093[gtid_5126]\
                \    let {lifted_lambda_res_5129 : i64} =\
                \      add64(2i64, eta_p_5128)\
                \    return {returns lifted_lambda_res_5129}\
                \  }"
          let res :: PrimExpTable
res = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Map VName (NameInfo MC) -> Stm MC -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Map VName (NameInfo MC)
scope Stm MC
stm) PrimExpTable
forall a. Monoid a => a
mempty
          let expected :: PrimExpTable
expected =
                [(VName, Maybe (PrimExp VName))] -> PrimExpTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  [ (VName
"defunc_0_map_res_5125", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    (VName
"gtid_5126", PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"gtid_5126" (IntType -> PrimType
IntType IntType
Int64))),
                    (VName
"eta_p_5128", Maybe (PrimExp VName)
forall a. Maybe a
Nothing),
                    ( VName
"lifted_lambda_res_5129",
                      PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just
                        ( BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp
                            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowWrap)
                            (PrimValue -> PrimExp VName
forall v. PrimValue -> PrimExp v
ValueExp (IntValue -> PrimValue
IntValue (Int64 -> IntValue
Int64Value Int64
2)))
                            (VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
"eta_p_5128" (IntType -> PrimType
IntType IntType
Int64))
                        )
                    )
                  ]
          PrimExpTable
res PrimExpTable -> PrimExpTable -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PrimExpTable
expected
        ]