content: module Futhark.Optimise.ArrayLayout.LayoutTests (tests) where
  
  import Data.Map.Strict qualified as M
  import Futhark.Analysis.AccessPattern
  import Futhark.Analysis.PrimExp
  import Futhark.FreshNames
  import Futhark.IR.GPU (GPU)
  import Futhark.IR.GPUTests ()
  import Futhark.Optimise.ArrayLayout.Layout
  import Language.Futhark.Core
  import Test.Tasty
  import Test.Tasty.HUnit
  
  tests :: TestTree
  tests =
    testGroup
      "Layout"
      [commonPermutationEliminatorsTests]
  
  commonPermutationEliminatorsTests :: TestTree
  commonPermutationEliminatorsTests =
    testGroup
      "commonPermutationEliminators"
      [permutationTests, nestTests, dimAccessTests, constIndexElimTests]
  
  permutationTests :: TestTree
  permutationTests =
    testGroup "Permutations" $
      do
        -- This isn't the way to test this, in reality we should provide realistic
        -- access patterns that might result in the given permutations.
        -- Luckily we only use the original access for one check atm.
        [ testCase (unwords [show perm, "->", show res]) $
            commonPermutationEliminators perm [] @?= res
          | (perm, res) <-
              [ ([0], True),
                ([1, 0], False),
                ([0, 1], True),
                ([0, 0], True),
                ([1, 1], True),
                ([1, 2, 0], False),
                ([2, 0, 1], False),
                ([0, 1, 2], True),
                ([1, 0, 2], True),
                ([2, 1, 0], True),
                ([2, 2, 0], True),
                ([2, 1, 1], True),
                ([1, 0, 1], True),
                ([0, 0, 0], True),
                ([0, 1, 2, 3, 4], True),
                ([1, 0, 2, 3, 4], True),
                ([2, 3, 0, 1, 4], True),
                ([3, 4, 2, 0, 1], True),
                ([2, 3, 4, 0, 1], False),
                ([1, 2, 3, 4, 0], False),
                ([3, 4, 0, 1, 2], False)
              ]
          ]
  
  nestTests :: TestTree
  nestTests = testGroup "Nests" $
    do
      let names = generateNames 2
      [ testCase (unwords [args, "->", show res]) $
          commonPermutationEliminators [1, 0] nest @?= res
        | (args, nest, res) <-
            [ ("[]", [], False),
              ("[CondBodyName]", [CondBodyName] <*> names, False),
              ("[SegOpName]", [SegOpName . SegmentedMap] <*> names, True),
              ("[LoopBodyName]", [LoopBodyName] <*> names, False),
              ("[SegOpName, CondBodyName]", [SegOpName . SegmentedMap, CondBodyName] <*> names, True),
              ("[CondBodyName, LoopBodyName]", [CondBodyName, LoopBodyName] <*> names, False)
            ]
        ]
  
  dimAccessTests :: TestTree
  dimAccessTests = testGroup "DimAccesses" [] -- TODO: Write tests for the part of commonPermutationEliminators that checks the complexity of the DimAccesses.
  
  constIndexElimTests :: TestTree
  constIndexElimTests =
    testGroup
      "constIndexElimTests"
      [ testCase "gpu eliminates indexes with constant in any dim" $ do
          let primExpTable =
                M.fromList
                  [ ("gtid_4", Just (LeafExp "n_4" (IntType Int64))),
                    ("i_5", Just (LeafExp "n_4" (IntType Int64)))
                  ]
          layoutTableFromIndexTable primExpTable accessTableGPU @?= mempty,
        testCase "gpu ignores when not last" $ do
          let primExpTable =
                M.fromList
                  [ ("gtid_4", Just (LeafExp "gtid_4" (IntType Int64))),
                    ("gtid_5", Just (LeafExp "gtid_5" (IntType Int64))),
                    ("i_6", Just (LeafExp "i_6" (IntType Int64)))
                  ]
          layoutTableFromIndexTable primExpTable accessTableGPUrev
            @?= M.fromList
              [ ( SegmentedMap "mapres_1",
                  M.fromList
                    [ ( ("a_2", [], [0, 1, 2, 3]),
                        M.fromList [("A_3", [2, 3, 0, 1])]
                      )
                    ]
                )
              ]
      ]
    where
      accessTableGPU :: IndexTable GPU
      accessTableGPU =
        singleAccess
          [ singleParAccess 0 "gtid_4",
            DimAccess mempty Nothing,
            singleSeqAccess 1 "i_5"
          ]
  
      accessTableGPUrev :: IndexTable GPU
      accessTableGPUrev =
        singleAccess
          [ singleParAccess 1 "gtid_4",
            singleParAccess 2 "gtid_5",
            singleSeqAccess 0 "i_5",
            singleSeqAccess 2 "gtid_4"
          ]
  
  singleAccess :: [DimAccess rep] -> IndexTable rep
  singleAccess dims =
    M.fromList
      [ ( sgOp,
          M.fromList
            [ ( ("A_2", [], [0, 1, 2, 3]),
                M.fromList
                  [ ( "a_3",
                      dims
                    )
                  ]
              )
            ]
        )
      ]
    where
      sgOp = SegmentedMap {vnameFromSegOp = "mapres_1"}
  
  singleParAccess :: Int -> VName -> DimAccess rep
  singleParAccess level name =
    DimAccess
      (M.singleton name $ Dependency level ThreadID)
      (Just name)
  
  singleSeqAccess :: Int -> VName -> DimAccess rep
  singleSeqAccess level name =
    DimAccess
      (M.singleton name $ Dependency level LoopVar)
      (Just name)
  
  generateNames :: Int -> [VName]
  generateNames count = do
    let (name, source) = newName blankNameSource "i_0"
    fst $ foldl f ([name], source) [1 .. count - 1]
    where
      f (names, source) _ = do
        let (name, source') = newName source (last names)
        (names ++ [name], source')
  