content: {-# LANGUAGE ApplicativeDo #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# OPTIONS_GHC -Wno-name-shadowing -Wno-incomplete-uni-patterns #-}
  
  -- |
  -- Tests for mutable vectors.
  --
  -- See the testing framework explained in Test.Data.Mutable.Set.
  --
  -- The combination of axioms and homomorphisms provided functionally specify
  -- the behavior of vectors.
  --
  -- Remarks:
  --  * We don't test for failure on out-of-bound access
  --  * We don't test the empty constructor
  module Test.Data.Mutable.Vector
    ( mutVecTests,
    )
  where
  
  import qualified Data.Functor.Linear as Data
  import qualified Data.List as List
  import Data.Maybe (mapMaybe)
  import Data.Ord.Linear as Linear hiding (Eq (..))
  import Data.Unrestricted.Linear
  import qualified Data.Vector as ImmutableVector
  import qualified Data.Vector.Mutable.Linear as Vector
  import Hedgehog
  import qualified Hedgehog.Gen as Gen
  import qualified Hedgehog.Range as Range
  import qualified Prelude.Linear as Linear hiding ((>))
  import Test.Tasty (TestTree, testGroup)
  import Test.Tasty.Hedgehog (testPropertyNamed)
  
  -- # Exported Tests
  --------------------------------------------------------------------------------
  
  mutVecTests :: TestTree
  mutVecTests = testGroup "Mutable vector tests" group
  
  group :: [TestTree]
  group =
    -- All tests for exprs of the form (read (const ...) i)
    [ testPropertyNamed "∀ s,i,x. read (constant s x) i = x" "readConst" readConst,
      testPropertyNamed "∀ a,i,x. read (write a i x) i = x " "readWrite1" readWrite1,
      testPropertyNamed "∀ a,i,j/=i,x. read (write a j x) i = read a i" "readWrite2" readWrite2,
      testPropertyNamed "∀ a,x,(i < len a). read (push a x) i = read a i" "readPush1" readPush1,
      testPropertyNamed "∀ a,x. read (push a x) (len a) = x" "readPush2" readPush2,
      -- All tests for exprs of the form (length (const ...))
      testPropertyNamed "∀ s,x. len (constant s x) = s" "lenConst" lenConst,
      testPropertyNamed "∀ a,i,x. len (write a i x) = len a" "lenWrite" lenWrite,
      testPropertyNamed "∀ a,x. len (push a x) = 1 + len a" "lenPush" lenPush,
      -- Tests against a reference implementation
      testPropertyNamed
        "write ix a v = (\\l -> take ix l ++ [a] ++ drop (ix+1) l) . toList"
        "refWrite"
        refWrite,
      testPropertyNamed "fst $ modify f ix v = snd $ f ((toList v) !! ix)" "refModify1" refModify1,
      testPropertyNamed
        "snd (modify f i v) = write (toList v) i (fst (f ((toList v) !! i))))"
        "refModify2"
        refModify2,
      testPropertyNamed "toList . push x = snoc x . toList" "refPush" refPush,
      testPropertyNamed "toList . pop = init . toList" "refPop" refPop,
      testPropertyNamed "read ix v = (toList v) !! ix" "refRead" refRead,
      testPropertyNamed "size = length . toList" "refSize" refSize,
      testPropertyNamed "toList . shrinkToFit = toList" "refShrinkToFit" refShrinkToFit,
      testPropertyNamed "pop . push _ = id" "refPopPush" refPopPush,
      testPropertyNamed "push . pop = id" "refPushPop" refPushPop,
      testPropertyNamed "slice s n = take s . drop n" "refSlice" refSlice,
      testPropertyNamed "toList . fromList = id" "refToListFromList" refToListFromList,
      testPropertyNamed "toList can be implemented with repeated pops" "refToListViaPop" refToListViaPop,
      testPropertyNamed "fromList can be implemented with repeated pushes" "refFromListViaPush" refFromListViaPush,
      testPropertyNamed "toList works with extra capacity" "refToListWithExtraCapacity" refToListWithExtraCapacity,
      testPropertyNamed "fromList xs <> fromList ys = fromList (xs <> ys)" "refMappend" refMappend,
      testPropertyNamed "mapMaybe f (fromList xs) = fromList (mapMaybe f xs)" "refMapMaybe" refMapMaybe,
      testPropertyNamed "filter f (fromList xs) = fromList (filter f xs)" "refFilter" refFilter,
      testPropertyNamed "f <$> fromList xs == fromList (f <$> xs)" "refFmap" refFmap,
      testPropertyNamed "toList . freeze . fromList = id" "refFreeze" refFreeze,
      -- Regression tests
      testPropertyNamed "push on an empty vector should succeed" "snocOnEmptyVector" snocOnEmptyVector,
      testPropertyNamed "do not reorder reads and writes" "readAndWriteTest" readAndWriteTest
    ]
  
  -- # Internal Library
  --------------------------------------------------------------------------------
  
  type VectorTester = Vector.Vector Int %1 -> Ur (TestT IO ())
  
  nonEmptyList :: Gen [Int]
  nonEmptyList = Gen.list (Range.linear 1 1000) val
  
  list :: Gen [Int]
  list = Gen.list (Range.linear 0 1000) val
  
  val :: Gen Int
  val = Gen.int (Range.linear (-1000) 1000)
  
  compInts ::
    Ur Int %1 ->
    Ur Int %1 ->
    Ur (TestT IO ())
  compInts (Ur x) (Ur y) = Ur (x === y)
  
  -- XXX: This is a terrible name
  getFst :: (Consumable b) => (a, b) %1 -> a
  getFst (a, b) = lseq b a
  
  getSnd :: (Consumable a) => (a, b) %1 -> b
  getSnd (a, b) = lseq a b
  
  -- # Tests
  --------------------------------------------------------------------------------
  
  snocOnEmptyVector :: Property
  snocOnEmptyVector = withTests 1 . property $ do
    let Ur actual =
          Vector.empty Linear.$
            \vec ->
              Vector.push (42 :: Int) vec
                Linear.& Vector.get 0
                Linear.& getFst
    actual === 42
  
  -- | Constant should give us a constant vector.
  readConst :: Property
  readConst = property $ do
    size <- forAll $ Gen.int $ Range.linear 1 1000
    v <- forAll val
    ix <- forAll $ Gen.element [0 .. size - 1]
    test $ unur Linear.$ Vector.constant size v (readConstTest ix v)
  
  readConstTest :: Int -> Int -> VectorTester
  readConstTest ix val vec = compInts (getFst (Vector.read vec ix)) (move val)
  
  readWrite1 :: Property
  readWrite1 = property $ do
    l <- forAll nonEmptyList
    let size = length l
    ix <- forAll $ Gen.element [0 .. size - 1]
    v <- forAll val
    let tester = readWrite1Test ix v
    test $ unur Linear.$ Vector.fromList l tester
  
  readWrite1Test :: Int -> Int -> VectorTester
  readWrite1Test ix val vec =
    compInts (move val) (getFst Linear.$ Vector.read (Vector.write vec ix val) ix)
  
  readWrite2 :: Property
  readWrite2 = property $ do
    let list = Gen.list (Range.linearFrom 2 2 1000) val
    l <- forAll list
    let size = length l
    ix <- forAll $ Gen.element [0 .. size - 1]
    jx <- forAll $ Gen.element [z | z <- [0 .. size - 1], z /= ix]
    v <- forAll val
    let tester = readWrite2Test ix jx v
    test $ unur Linear.$ Vector.fromList l tester
  
  readWrite2Test :: Int -> Int -> Int -> VectorTester
  readWrite2Test ix jx val vec = fromRead (Vector.read vec ix)
    where
      fromRead :: (Ur Int, Vector.Vector Int) %1 -> Ur (TestT IO ())
      fromRead (val1, vec) =
        compInts
          val1
          (getFst Linear.$ Vector.read (Vector.write vec jx val) ix)
  
  readPush1 :: Property
  readPush1 = property $ do
    l <- forAll nonEmptyList
    let size = length l
    v <- forAll val
    ix <- forAll $ Gen.element [0 .. size - 1]
    let tester = readPush1Test v ix
    test $ unur Linear.$ Vector.fromList l tester
  
  readPush1Test :: Int -> Int -> VectorTester
  readPush1Test val ix vec = fromRead (Vector.read vec ix)
    where
      fromRead :: (Ur Int, Vector.Vector Int) %1 -> Ur (TestT IO ())
      fromRead (val', vec) =
        compInts (getFst (Vector.get ix (Vector.push val vec))) val'
  
  readPush2 :: Property
  readPush2 = property $ do
    l <- forAll list
    v <- forAll val
    let tester = readPush2Test v
    test $ unur Linear.$ Vector.fromList l tester
  
  readPush2Test :: Int -> VectorTester
  readPush2Test val vec = fromLen (Vector.size vec)
    where
      fromLen ::
        (Ur Int, Vector.Vector Int) %1 ->
        Ur (TestT IO ())
      fromLen (Ur len, vec) =
        compInts (getFst (Vector.get len (Vector.push val vec))) (move val)
  
  lenConst :: Property
  lenConst = property $ do
    size <- forAll $ Gen.int $ Range.linear 1 1000
    v <- forAll val
    test $ unur Linear.$ Vector.constant size v (lenConstTest size)
  
  lenConstTest :: Int -> VectorTester
  lenConstTest size vec =
    compInts (move size) (getFst Linear.$ Vector.size vec)
  
  lenWrite :: Property
  lenWrite = property $ do
    l <- forAll nonEmptyList
    let size = length l
    v <- forAll val
    ix <- forAll $ Gen.element [0 .. size - 1]
    let tester = lenWriteTest size v ix
    test $ unur Linear.$ Vector.fromList l tester
  
  lenWriteTest :: Int -> Int -> Int -> VectorTester
  lenWriteTest size val ix vec =
    compInts
      (move size)
      (getFst Linear.$ Vector.size (Vector.write vec ix val))
  
  lenPush :: Property
  lenPush = property $ do
    l <- forAll list
    v <- forAll val
    let tester = lenPushTest v
    test $ unur Linear.$ Vector.fromList l tester
  
  lenPushTest :: Int -> VectorTester
  lenPushTest val vec = fromLen Linear.$ Vector.size vec
    where
      fromLen ::
        (Ur Int, Vector.Vector Int) %1 ->
        Ur (TestT IO ())
      fromLen (Ur len, vec) =
        compInts (move (len + 1)) (getFst (Vector.size (Vector.push val vec)))
  
  refWrite :: Property
  refWrite = property $ do
    l <- forAll nonEmptyList
    ix <- forAll $ Gen.element [0 .. (length l - 1)]
    v <- forAll val
    let l' = listWrite ix v l
    l' === unur (Vector.fromList l (Vector.toList Linear.. Vector.set ix v))
    where
      listWrite :: Int -> a -> [a] -> [a]
      listWrite n _ _ | n Prelude.< 0 = error "Index negative"
      listWrite _ _ [] = error "Index too big"
      listWrite 0 a (_ : xs) = a : xs
      listWrite n a (x : xs) = x : listWrite (n - 1) a xs
  
  refModify1 :: Property
  refModify1 = property $ do
    l <- forAll nonEmptyList
    let f x = (mod x 5, (mod x 5) Prelude.< 3)
    ix <- forAll $ Gen.element [0 .. (length l - 1)]
    snd (f (l !! ix)) === unur (Vector.fromList l (getFst Linear.. Vector.modify f ix))
  
  refModify2 :: Property
  refModify2 = property $ do
    l <- forAll nonEmptyList
    let f x = 3 * x * x - 2 * x + 4
    ix <- forAll $ Gen.element [0 .. (length l - 1)]
    let l' = listMod ix f l
    l' === unur (Vector.fromList l (Vector.toList Linear.. Vector.modify_ f ix))
    where
      listMod :: Int -> (a -> a) -> [a] -> [a]
      listMod n _ _ | n Prelude.< 0 = error "Index negative"
      listMod _ _ [] = error "Index too big"
      listMod 0 f (x : xs) = f x : xs
      listMod n f (x : xs) = x : listMod (n - 1) f xs
  
  refPush :: Property
  refPush = property $ do
    l <- forAll list
    v <- forAll val
    let l' = l ++ [v]
    l' === unur (Vector.fromList l (Vector.toList Linear.. Vector.push v))
  
  refPop :: Property
  refPop = property $ do
    l <- forAll nonEmptyList
    let v = Vector.fromList l (Vector.toList Linear.. getSnd Linear.. Vector.pop)
    List.init l === unur v
  
  refRead :: Property
  refRead = property $ do
    l <- forAll nonEmptyList
    ix <- forAll $ Gen.element [0 .. (length l - 1)]
    let value = l List.!! ix
    value === unur (Vector.fromList l (getFst Linear.. Vector.get ix))
  
  refSize :: Property
  refSize = property $ do
    l <- forAll list
    length l === unur (Vector.fromList l (getFst Linear.. Vector.size))
  
  refShrinkToFit :: Property
  refShrinkToFit = property $ do
    l <- forAll list
    l === unur (Vector.fromList l (Vector.toList Linear.. Vector.shrinkToFit))
  
  refToListFromList :: Property
  refToListFromList = property $ do
    xs <- forAll list
    let Ur actual = Vector.fromList xs Vector.toList
    xs === actual
  
  refToListWithExtraCapacity :: Property
  refToListWithExtraCapacity = property $ do
    xs <- forAll list
    let val = 12
        expected = xs ++ [val]
        Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.push val vec -- This will cause the vector to grow.
              Linear.& Vector.toList
    expected === actual
  
  refPopPush :: Property
  refPopPush = property $ do
    xs <- forAll list
    let Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.push (error "not used") vec
              Linear.& Vector.pop
              Linear.& \(Ur _, vec) ->
                Vector.toList vec
    xs === actual
  
  refPushPop :: Property
  refPushPop = property $ do
    xs <- forAll nonEmptyList
    let Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.pop vec
              Linear.& \(Ur (Just a), vec) ->
                Vector.push a vec
                  Linear.& Vector.toList
    xs === actual
  
  refToListViaPop :: Property
  refToListViaPop = property $ do
    xs <- forAll list
    let Ur actual =
          Vector.fromList xs (popAll [])
    xs === actual
    where
      popAll :: [a] -> Vector.Vector a %1 -> Ur [a]
      popAll acc vec =
        case Vector.pop vec of
          (Ur Nothing, vec') -> vec' `lseq` Ur acc
          (Ur (Just x), vec') -> popAll (x : acc) vec'
  
  refFromListViaPush :: Property
  refFromListViaPush = property $ do
    xs <- forAll list
    let Ur actual =
          Vector.empty Linear.$
            Vector.toList Linear.. pushAll xs
    xs === actual
    where
      pushAll :: [a] -> Vector.Vector a %1 -> Vector.Vector a
      pushAll [] vec = vec
      pushAll (x : xs) vec = Vector.push x vec Linear.& pushAll xs
  
  refSlice :: Property
  refSlice = property $ do
    xs <- forAll list
    s <- forAll $ Gen.int (Range.linear 0 (length xs))
    n <- forAll $ Gen.int (Range.linear 0 (length xs - s))
    let expected = take n . drop s $ xs
        Ur actual =
          Vector.fromList xs Linear.$ \arr ->
            Vector.slice s n arr
              Linear.& Vector.toList
    expected === actual
  
  refMappend :: Property
  refMappend = property $ do
    xs <- forAll list
    ys <- forAll list
    let expected = xs <> ys
        Ur actual =
          Vector.fromList xs Linear.$ \vx ->
            Vector.fromList ys Linear.$ \vy ->
              Vector.toList (vx Linear.<> vy)
    expected === actual
  
  refFmap :: Property
  refFmap = property $ do
    xs <- forAll list
    let -- An arbitrary function
        f :: Int %1 -> Bool
        f = (Linear.> 0)
        expected = map (Linear.forget f) xs
        Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.toList (f Data.<$> vec)
    expected === actual
  
  refMapMaybe :: Property
  refMapMaybe = property $ do
    xs <- forAll list
    let -- An arbitrary function
        f :: Int -> Maybe Bool
        f = (\a -> if a Prelude.< 0 then Nothing else Just (a `mod` 2 == 0))
        expected = mapMaybe f xs
        Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.mapMaybe vec f
              Linear.& Vector.toList
    expected === actual
  
  refFilter :: Property
  refFilter = property $ do
    xs <- forAll list
    let -- An arbitrary function
        f :: Int -> Bool
        f = (Prelude.< 0)
        expected = filter f xs
        Ur actual =
          Vector.fromList xs Linear.$ \vec ->
            Vector.filter vec f
              Linear.& Vector.toList
    expected === actual
  
  refFreeze :: Property
  refFreeze = property $ do
    xs <- forAll list
  
    -- Add a new element at the end of the vector
    -- to force resizing, to test the case where
    -- sz < cap.
    shouldAppend <- forAll Gen.bool
  
    let expected =
          if shouldAppend
            then xs ++ [12]
            else xs
  
        Ur actual = Vector.fromList xs Linear.$ \vec ->
          ( if shouldAppend
              then Vector.push 12 vec
              else vec
          )
            Linear.& Vector.freeze
    expected === ImmutableVector.toList actual
  
  -- https://github.com/tweag/linear-base/pull/135
  readAndWriteTest :: Property
  readAndWriteTest =
    withTests 1 . property $
      unur (Vector.fromList "a" test) === 'a'
    where
      test :: Vector.Vector Char %1 -> Ur Char
      test vec =
        Vector.read vec 0 Linear.& \(before, vec') ->
          Vector.write vec' 0 'b' Linear.& \vec'' ->
            vec'' `Linear.lseq` before
  