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 = testGroup "AnalyzePrim" [stmToPrimExpsTests]

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

stmToPrimExpsTestsGPU :: TestTree
stmToPrimExpsTestsGPU =
  testGroup
    "GPU"
    $ do
      let scope =
            M.fromList
              [ ("n_5142", FParamName "i64"),
                ("m_5143", FParamName "i64"),
                ("xss_5144", FParamName "[n_5142][m_5143]i64"),
                ("segmap_group_size_5201", LetName "i64"),
                ("segmap_usable_groups_5202", LetName "i64"),
                ("defunc_0_map_res_5203", LetName "[n_5142]i64"),
                ("defunc_0_f_res_5207", LetName "i64"),
                ("i_5208", IndexName Int64),
                ("acc_5209", FParamName "i64"),
                ("b_5210", LetName "i64"),
                ("defunc_0_f_res_5211", LetName "i64")
              ]
      [ testCase "BinOp" $ do
          let stm = "let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"
          let res = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ( "defunc_0_f_res_5211",
                      Just
                        ( BinOpExp
                            (Add Int64 OverflowWrap)
                            (LeafExp "acc_5209" (IntType Int64))
                            (LeafExp "b_5210" (IntType Int64))
                        )
                    )
                  ]
          res @?= expected,
        testCase "Index" $ do
          let stm = "let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"
          let res = execState (stmToPrimExps scope stm) mempty
          let expected = M.fromList [("b_5210", Nothing)]
          res @?= expected,
        testCase "Loop" $ do
          let stm = "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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_f_res_5207", Nothing),
                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),
                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64)))
                  ]
          res @?= expected,
        testCase "Loop body" $ do
          let stm = "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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_f_res_5207", Nothing),
                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),
                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64))),
                    ("b_5210", Nothing),
                    ( "defunc_0_f_res_5211",
                      Just
                        ( BinOpExp
                            (Add Int64 OverflowWrap)
                            (LeafExp "acc_5209" (IntType Int64))
                            (LeafExp "b_5210" (IntType Int64))
                        )
                    )
                  ]
          res @?= expected,
        testCase "SegMap" $
          do
            let stm =
                  "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 = execState (stmToPrimExps scope stm) mempty
            let expected =
                  M.fromList
                    [ ("defunc_0_map_res_5125", Nothing),
                      ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64)))
                    ]
            res @?= expected,
        testCase "SegMap body" $
          do
            let stm :: Stm GPU
                stm =
                  "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 = execState (stmToPrimExps scope stm) mempty
            let expected =
                  M.fromList
                    [ ("defunc_0_map_res_5125", Nothing),
                      ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64))),
                      ("eta_p_5128", Nothing),
                      ( "lifted_lambda_res_5129",
                        Just
                          ( BinOpExp
                              (Add Int64 OverflowWrap)
                              (ValueExp (IntValue (Int64Value 2)))
                              (LeafExp "eta_p_5128" (IntType Int64))
                          )
                      )
                    ]
            res @?= expected
        ]

stmToPrimExpsTestsMC :: TestTree
stmToPrimExpsTestsMC =
  testGroup
    "MC"
    $ do
      let scope =
            M.fromList
              [ ("n_5142", FParamName "i64"),
                ("m_5143", FParamName "i64"),
                ("xss_5144", FParamName "[n_5142][5143]i64"),
                ("segmap_group_size_5201", LetName "i64"),
                ("segmap_usable_groups_5202", LetName "i64"),
                ("defunc_0_map_res_5203", LetName "[n_5142]i64"),
                ("defunc_0_f_res_5207", LetName "i64"),
                ("i_5208", IndexName Int64),
                ("acc_5209", FParamName "i64"),
                ("b_5210", LetName "i64"),
                ("defunc_0_f_res_5211", LetName "i64")
              ]
      [ testCase "BinOp" $ do
          let stm = "let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"
          let res = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ( "defunc_0_f_res_5211",
                      Just
                        ( BinOpExp
                            (Add Int64 OverflowWrap)
                            (LeafExp "acc_5209" (IntType Int64))
                            (LeafExp "b_5210" (IntType Int64))
                        )
                    )
                  ]
          res @?= expected,
        testCase "Index" $ do
          let stm = "let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"
          let res = execState (stmToPrimExps scope stm) mempty
          let expected = M.fromList [("b_5210", Nothing)]
          res @?= expected,
        testCase "Loop" $ do
          let stm = "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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_f_res_5207", Nothing),
                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),
                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64)))
                  ]
          res @?= expected,
        testCase "Loop body" $ do
          let stm =
                "\
                \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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_f_res_5207", Nothing),
                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),
                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64))),
                    ("b_5210", Nothing),
                    ( "defunc_0_f_res_5211",
                      Just
                        ( BinOpExp
                            (Add Int64 OverflowWrap)
                            (LeafExp "acc_5209" (IntType Int64))
                            (LeafExp "b_5210" (IntType Int64))
                        )
                    )
                  ]
          res @?= expected,
        testCase "SegMap" $ do
          let stm =
                "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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_map_res_5125", Nothing),
                    ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64)))
                  ]
          res @?= expected,
        testCase "SegMap body" $ do
          let stm :: Stm MC
              stm =
                "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 = execState (stmToPrimExps scope stm) mempty
          let expected =
                M.fromList
                  [ ("defunc_0_map_res_5125", Nothing),
                    ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64))),
                    ("eta_p_5128", Nothing),
                    ( "lifted_lambda_res_5129",
                      Just
                        ( BinOpExp
                            (Add Int64 OverflowWrap)
                            (ValueExp (IntValue (Int64Value 2)))
                            (LeafExp "eta_p_5128" (IntType Int64))
                        )
                    )
                  ]
          res @?= expected
        ]