content: module Futhark.IR.Prop.RearrangeTests (tests) where
  
  import Control.Applicative
  import Futhark.IR.Prop.Rearrange
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.QuickCheck
  import Prelude
  
  tests :: TestTree
  tests =
    testGroup "RearrangeTests" $
      isMapTransposeTests
        ++ [isMapTransposeProp]
  
  isMapTransposeTests :: [TestTree]
  isMapTransposeTests =
    [ testCase (unwords ["isMapTranspose", show perm, "==", show dres]) $
        isMapTranspose perm @?= dres
      | (perm, dres) <-
          [ ([0, 1, 4, 5, 2, 3], Just (2, 2, 2)),
            ([1, 0, 4, 5, 2, 3], Nothing),
            ([1, 0], Just (0, 1, 1)),
            ([0, 2, 1], Just (1, 1, 1)),
            ([0, 1, 2], Nothing),
            ([1, 0, 2], Nothing)
          ]
    ]
  
  newtype Permutation = Permutation [Int]
    deriving (Eq, Ord, Show)
  
  instance Arbitrary Permutation where
    arbitrary = do
      Positive n <- arbitrary
      Permutation <$> shuffle [0 .. n - 1]
  
  isMapTransposeProp :: TestTree
  isMapTransposeProp = testProperty "isMapTranspose corresponds to a map of transpose" prop
    where
      prop :: Permutation -> Bool
      prop (Permutation perm) =
        case isMapTranspose perm of
          Nothing -> True
          Just (r1, r2, r3) ->
            and
              [ r1 >= 0,
                r2 > 0,
                r3 > 0,
                r1 + r2 + r3 == length perm,
                let (mapped, notmapped) = splitAt r1 perm
                    (pretrans, posttrans) = splitAt r2 notmapped
                 in mapped ++ posttrans ++ pretrans == [0 .. length perm - 1]
              ]
  