content: {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module Futhark.IR.Prop.ReshapeTests
    ( tests,
    )
  where
  
  import Data.List qualified as L
  import Futhark.IR.Prop.Constants
  import Futhark.IR.Prop.Reshape
  import Futhark.IR.Syntax
  import Futhark.IR.SyntaxTests ()
  import Test.Tasty
  import Test.Tasty.HUnit
  
  intShape :: [Int] -> Shape
  intShape = Shape . map (intConst Int32 . toInteger)
  
  reshapeOuterTests :: [TestTree]
  reshapeOuterTests =
    [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $
        reshapeOuter (intShape sc) n (intShape shape) @?= intShape sc_res
      | (sc, n, shape, sc_res) <-
          [ ([1], 1, [4, 3], [1, 3]),
            ([1], 2, [4, 3], [1]),
            ([2, 2], 1, [4, 3], [2, 2, 3]),
            ([2, 2], 2, [4, 3], [2, 2])
          ]
    ]
  
  reshapeInnerTests :: [TestTree]
  reshapeInnerTests =
    [ testCase (unwords ["reshapeInner", show sc, show n, show shape, "==", show sc_res]) $
        reshapeInner (intShape sc) n (intShape shape) @?= intShape sc_res
      | (sc, n, shape, sc_res) <-
          [ ([1], 1, [4, 3], [4, 1]),
            ([1], 0, [4, 3], [1]),
            ([2, 2], 1, [4, 3], [4, 2, 2]),
            ([2, 2], 0, [4, 3], [2, 2])
          ]
    ]
  
  dimFlatten :: Int -> Int -> d -> DimSplice d
  dimFlatten i k w = DimSplice i k (Shape [w])
  
  dimUnflatten :: Int -> [d] -> DimSplice d
  dimUnflatten i ws = DimSplice i 1 (Shape ws)
  
  dimCoerce :: Int -> d -> DimSplice d
  dimCoerce i w = DimSplice i 1 (Shape [w])
  
  dimSplice :: Int -> Int -> [d] -> DimSplice d
  dimSplice i n s = DimSplice i n $ Shape s
  
  flipReshapeRearrangeTests :: [TestTree]
  flipReshapeRearrangeTests =
    [ testCase
        ( unwords
            [ "flipReshapeRearrange",
              show v0_shape,
              show v1_shape,
              show perm
            ]
        )
        $ flipReshapeRearrange v0_shape v1_shape perm @?= res
      | (v0_shape :: [String], v1_shape, perm, res) <-
          [ ( ["A", "B", "C"],
              ["A", "BC"],
              [1, 0],
              Just [1, 2, 0]
            ),
            ( ["A", "B", "C", "D"],
              ["A", "BCD"],
              [1, 0],
              Just [1, 2, 3, 0]
            ),
            ( ["A"],
              ["B", "C"],
              [1, 0],
              Nothing
            ),
            ( ["A", "B", "C"],
              ["AB", "C"],
              [1, 0],
              Just [2, 0, 1]
            ),
            ( ["A", "B", "C", "D"],
              ["ABC", "D"],
              [1, 0],
              Just [3, 0, 1, 2]
            )
          ]
    ]
  
  flipRearrangeReshapeTests :: [TestTree]
  flipRearrangeReshapeTests =
    [ testCase
        ( unwords
            [ "flipRearrangeReshape",
              show perm,
              prettyStringOneLine newshape
            ]
        )
        $ flipRearrangeReshape perm newshape @?= res
      | (perm, newshape :: NewShape String, res) <-
          [ ( [1, 0],
              NewShape
                [dimUnflatten 1 ["B", "C"]]
                (Shape ["A", "B", "C"]),
              Just
                ( NewShape
                    [dimUnflatten 0 ["B", "C"]]
                    (Shape ["B", "C", "A"]),
                  [2, 0, 1]
                )
            ),
            ( [1, 0],
              NewShape
                [dimFlatten 0 2 "AB"]
                (Shape ["AB"]),
              Nothing
            )
          ]
    ]
  
  simplifyTests :: TestTree
  simplifyTests =
    testGroup
      "simplifyNewShape"
      [ testCase "Inverse flatten and unflatten - simple case" $
          lhs
            ["A", "B"]
            [dimFlatten 0 2 "AB", dimUnflatten 0 ["A", "B"]]
            @?= Just [],
        testCase "Non-inverse flatten and unflatten - simple case" $
          lhs
            ["A", "B"]
            [dimFlatten 0 2 "AB", dimUnflatten 0 ["C", "D"]]
            @?= Just [dimSplice 0 2 ["C", "D"]],
        testCase "Inverse flatten and unflatten - separated by coercion" $
          lhs
            ["A", "B"]
            [ dimFlatten 0 2 "AB",
              dimCoerce 0 "CD",
              dimUnflatten 0 ["C", "D"]
            ]
            @?= Just [dimSplice 0 2 ["C", "D"]],
        testCase "Two unflattens - simple case" $
          lhs
            ["ABC"]
            [dimUnflatten 0 ["A", "BC"], dimUnflatten 1 ["B", "C"]]
            @?= Just [dimUnflatten 0 ["A", "B", "C"]],
        testCase "Two unflattens with unchanged prefix" $
          lhs
            ["A", "B", "C", "D", "E"]
            [ DimSplice 3 2 $ Shape ["DE"],
              DimSplice 2 2 $ Shape ["CDE"]
            ]
            @?= Just [dimFlatten 2 3 "CDE"],
        testCase "Identity coerce" $
          lhs
            ["A", "B", "C"]
            [dimCoerce 1 "B", dimCoerce 2 "C"]
            @?= Just [],
        testCase "Identity coerce (multiple dimensions)" $
          lhs
            ["A", "B", "C"]
            [DimSplice 1 2 (Shape ["B", "C"])]
            @?= Just [],
        testCase "Identity coerce (with non-identity stuff afterwards)" $
          lhs
            ["B", "CD"]
            [dimCoerce 0 "B", dimUnflatten 1 ["C", "D"]]
            @?= Just [dimUnflatten 1 ["C", "D"]],
        testCase "Get rid of a coerce before an unflatten" $
          lhs
            ["CD"]
            [dimCoerce 0 "AB", dimUnflatten 0 ["A", "B"]]
            @?= Just [dimUnflatten 0 ["A", "B"]],
        testCase "Get rid of a coerce after a flatten" $
          lhs
            ["A", "B", "C"]
            [dimFlatten 0 2 "ABC", dimCoerce 0 "K"]
            @?= Just [dimFlatten 0 2 "K"],
        testCase "Flatten and unflatten (invariant suffix)" $
          lhs
            ["A", "B", "C"]
            [dimFlatten 0 3 "ABC", dimUnflatten 0 ["D", "E", "C"]]
            @?= Just [dimSplice 0 2 ["D", "E"]],
        testCase "Flatten and unflatten (invariant prefix)" $
          lhs
            ["A", "B", "C"]
            [dimFlatten 0 3 "ABC", dimUnflatten 0 ["A", "D", "E"]]
            @?= Just [dimSplice 1 2 ["D", "E"]],
        testCase "Invariant part of splice" $
          lhs
            ["A", "B", "C", "D"]
            [DimSplice 1 3 $ Shape ["BC", "D"]]
            @?= Just [DimSplice 1 2 $ Shape ["BC"]],
        testCase "Necessary coercion" $
          lhs
            ["A", "B"]
            [dimCoerce 0 "C", dimCoerce 1 "D"]
            @?= Nothing,
        testCase "Another necessary coercion" $
          lhs
            ["A", "B", "C"]
            [dimCoerce 0 "A'", dimCoerce 1 "A'", dimCoerce 2 "A'"]
            @?= Nothing,
        testCase "Long with redundancies" $
          lhs
            ["A", "B", "C", "D"]
            [ DimSplice 1 3 $ Shape ["BC", "D"],
              dimCoerce 1 "BC",
              dimCoerce 2 "D",
              dimFlatten 1 2 "BCD",
              dimFlatten 0 2 "ABCD"
            ]
            @?= Just [dimFlatten 0 4 "ABCD"]
      ]
    where
      lhs orig_shape ss =
        let res_shape :: ShapeBase String =
              L.foldl' applySplice (Shape orig_shape) ss
         in dimSplices
              <$> simplifyNewShape (Shape orig_shape) (NewShape ss res_shape)
  
  tests :: TestTree
  tests =
    testGroup "ReshapeTests" . mconcat $
      [ reshapeOuterTests,
        reshapeInnerTests,
        flipReshapeRearrangeTests,
        flipRearrangeReshapeTests,
        [simplifyTests]
      ]
  