content: {-# LANGUAGE OverloadedStrings #-}
  
  module Grisette.Lib.Data.FoldableTests (foldableFunctionTests) where
  
  import Control.Monad.Except
    ( ExceptT (ExceptT),
      MonadError (throwError),
      runExceptT,
    )
  import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
  import Grisette
    ( AsKey1,
      ITEOp (symIte),
      LogicalOp ((.&&), (.||)),
      SymBool,
      SymBranching (mrgIfPropagatedStrategy),
      SymEq ((./=), (.==)),
      SymInteger,
      SymOrd (symCompare, (.<=), (.>)),
      Union,
      mrgIf,
      symAll,
      symAny,
      symOr,
    )
  import Grisette.Lib.Control.Monad (mrgMzero, mrgReturn)
  import Grisette.Lib.Data.Foldable
    ( mrgFind,
      mrgFoldlM,
      mrgFoldrM,
      mrgForM_,
      mrgFor_,
      mrgMapM_,
      mrgMaximum,
      mrgMaximumBy,
      mrgMinimum,
      mrgMinimumBy,
      mrgMsum,
      mrgSequenceA_,
      mrgSequence_,
      mrgTraverse_,
      symAnd,
      symElem,
      symMaximum,
      symMaximumBy,
      symMinimum,
      symMinimumBy,
      symNotElem,
    )
  import Grisette.TestUtil.NoMerge (noMergeNotMerged, oneNotMerged)
  import Grisette.TestUtil.SymbolicAssertion ((.@?=))
  import Test.Framework
    ( Test,
      TestOptions' (topt_timeout),
      plusTestOptions,
      testGroup,
    )
  import Test.Framework.Providers.HUnit (testCase)
  import Test.HUnit ((@?=))
  
  foldableFunctionTests :: Test
  foldableFunctionTests =
    testGroup
      "Foldable"
      [ testCase "symElem" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symElem a [b, c]
          actual .@?= (a .== b .|| a .== c),
        testCase "mrgMaximum" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = mrgMaximum [[a], [b, a], [c, a, b]]
          let expected =
                mrgIf
                  (a .> b .&& a .> c)
                  (return [a])
                  (mrgIf (b .> c) (return [b, a]) (return [c, a, b])) ::
                  Union [SymInteger]
          actual .@?= expected,
        testCase "symMaximum" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symMaximum [a, b, c]
          let expected = symIte (a .> b .&& a .> c) a (symIte (b .> c) b c)
          actual .@?= expected,
        testCase "mrgMinimum" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = mrgMinimum [[a], [b, a], [c, a, b]]
          let expected =
                mrgIf
                  (a .<= b .&& a .<= c)
                  (return [a])
                  (mrgIf (b .<= c) (return [b, a]) (return [c, a, b])) ::
                  Union [SymInteger]
          actual .@?= expected,
        testCase "symMinimum" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symMinimum [a, b, c]
          let expected = symIte (a .<= b .&& a .<= c) a (symIte (b .<= c) b c)
          actual .@?= expected,
        plusTestOptions (mempty {topt_timeout = Just (Just 1000000)}) $
          testGroup
            "mrgFoldrM"
            [ testCase "semantics" $ do
                let actual =
                      mrgFoldrM
                        ( \(c, v) acc ->
                            mrgIfPropagatedStrategy
                              c
                              (return $ acc + v)
                              (return $ acc * v)
                        )
                        10
                        [("a", 2), ("b", 3)] ::
                        AsKey1 Union Integer
                let expected =
                      mrgIf
                        "b"
                        (mrgIf "a" (mrgReturn 15) (mrgReturn 26))
                        (mrgIf "a" (mrgReturn 32) (mrgReturn 60))
                actual @?= expected,
              testCase "merge intermediate" $ do
                let actual = mrgFoldrM (const $ const oneNotMerged) 1 [1 .. 1000]
                let expected = mrgReturn 1 :: AsKey1 Union Int
                actual @?= expected
            ],
        plusTestOptions (mempty {topt_timeout = Just (Just 1000000)}) $
          testGroup
            "mrgFoldlM"
            [ testCase "semantics" $ do
                let actual =
                      mrgFoldlM
                        ( \acc (c, v) ->
                            mrgIfPropagatedStrategy
                              c
                              (return $ acc + v)
                              (return $ acc * v)
                        )
                        10
                        [("a", 2), ("b", 3)] ::
                        AsKey1 Union Integer
                let expected =
                      mrgIf
                        "a"
                        (mrgIf "b" (mrgReturn 15) (mrgReturn 36))
                        (mrgIf "b" (mrgReturn 23) (mrgReturn 60))
                actual @?= expected,
              testCase "merge intermediate" $ do
                let actual = mrgFoldlM (const $ const oneNotMerged) 1 [1 .. 1000]
                let expected = mrgReturn 1 :: AsKey1 Union Int
                actual @?= expected
            ],
        plusTestOptions (mempty {topt_timeout = Just (Just 1000000)}) $
          testGroup "mrgTraverse_, mrgMapM_, mrgFor_, mrgForM_" $ do
            (name, func0, func1) <-
              [ ("mrgTraverse_", mrgTraverse_, mrgTraverse_),
                ("mrgMapM_", mrgMapM_, mrgMapM_),
                ("mrgFor_", flip mrgFor_, flip mrgFor_),
                ("mrgForM_", flip mrgForM_, flip mrgForM_)
              ]
            [ testGroup
                name
                [ testCase "semantics" $ do
                    let actual =
                          runExceptT
                            ( func0
                                ( \(c, x) ->
                                    ExceptT $
                                      mrgIfPropagatedStrategy
                                        c
                                        (return $ Left x)
                                        (return $ Right c)
                                )
                                [("a", 3), ("b", 2)] ::
                                ExceptT Integer (AsKey1 Union) ()
                            )
                    let expected = runExceptT $ do
                          _ <- mrgIf "a" (throwError 3) (return ())
                          _ <- mrgIf "b" (throwError 2) (return ())
                          mrgReturn ()
                    actual @?= expected,
                  testCase "discard and merge intermediate" $ do
                    let actual = func1 (const noMergeNotMerged) [1 .. 1000]
                    let expected = mrgReturn () :: AsKey1 Union ()
                    actual @?= expected
                ]
              ],
        plusTestOptions (mempty {topt_timeout = Just (Just 1000000)}) $
          testGroup "mrgSequence_, mrgSequenceA_" $ do
            (name, func0, func1) <-
              [ ("mrgSequence_", mrgSequence_, mrgSequence_),
                ("mrgSequenceA_", mrgSequenceA_, mrgSequenceA_)
              ]
            [ testGroup
                name
                [ testCase "semantics" $ do
                    let actual =
                          runExceptT
                            ( func0 $
                                ( \(c, x) ->
                                    ExceptT $
                                      mrgIfPropagatedStrategy
                                        c
                                        (return $ Left x)
                                        (return $ Right c)
                                )
                                  <$> [("a", 3), ("b", 2)] ::
                                ExceptT Integer (AsKey1 Union) ()
                            )
                    let expected = runExceptT $ do
                          _ <- mrgIf "a" (throwError 3) (return ())
                          _ <- mrgIf "b" (throwError 2) (return ())
                          mrgReturn ()
                    actual @?= expected,
                  testCase "discard and merge intermediate" $ do
                    let actual = func1 (replicate 1000 noMergeNotMerged)
                    let expected = mrgReturn () :: AsKey1 Union ()
                    actual @?= expected
                ]
              ],
        plusTestOptions (mempty {topt_timeout = Just (Just 1000000)}) $
          testGroup
            "mrgMsum"
            [ testCase "merge" $ do
                let none =
                      MaybeT $
                        mrgIfPropagatedStrategy "a" (return Nothing) (return Nothing)
                let expected =
                      MaybeT (mrgReturn Nothing) ::
                        MaybeT Union (Maybe Int)
                mrgMsum (replicate 100 none) @?= expected,
              testCase "semantics" $ do
                (mrgMsum [mrgMzero, mrgMzero] :: MaybeT Union Integer)
                  @?= mrgMzero
                (mrgMsum [mrgReturn 1, mrgMzero] :: MaybeT Union Integer)
                  @?= mrgReturn 1
                (mrgMsum [mrgMzero, mrgReturn 1] :: MaybeT Union Integer)
                  @?= mrgReturn 1
                (mrgMsum [mrgReturn 2, mrgReturn 1] :: MaybeT Union Integer)
                  @?= mrgReturn 2
            ],
        testCase "symAnd" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymBool]
          symAnd [a, b, c] .@?= (a .&& b .&& c),
        testCase "symOr" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymBool]
          symOr [a, b, c] .@?= (a .|| b .|| c),
        testCase "symAny" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          symAny (.== 0) [a, b, c] .@?= (a .== 0 .|| b .== 0 .|| c .== 0),
        testCase "symAll" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          symAll (.== 0) [a, b, c] .@?= (a .== 0 .&& b .== 0 .&& c .== 0),
        testCase "mrgMaximumBy" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = mrgMaximumBy symCompare [[a], [b, a], [c, a, b]]
          let expected =
                mrgIf
                  (a .> b .&& a .> c)
                  (return [a])
                  (mrgIf (b .> c) (return [b, a]) (return [c, a, b])) ::
                  Union [SymInteger]
          actual .@?= expected,
        testCase "symMaximumBy" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symMaximumBy symCompare [a, b, c]
          let expected = symIte (a .> b .&& a .> c) a (symIte (b .> c) b c)
          actual .@?= expected,
        testCase "mrgMinimumBy" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = mrgMinimumBy symCompare [[a], [b, a], [c, a, b]]
          let expected =
                mrgIf
                  (a .<= b .&& a .<= c)
                  (return [a])
                  (mrgIf (b .<= c) (return [b, a]) (return [c, a, b])) ::
                  Union [SymInteger]
          actual .@?= expected,
        testCase "symMinimumBy" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symMinimumBy symCompare [a, b, c]
          let expected = symIte (a .<= b .&& a .<= c) a (symIte (b .<= c) b c)
          actual .@?= expected,
        testCase "symNotElem" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = symNotElem a [b, c]
          actual .@?= (a ./= b .&& a ./= c),
        testCase "mrgFind" $ do
          let [a, b, c] = ["a", "b", "c"] :: [SymInteger]
          let actual = mrgFind (.== 0) [a, b, c] :: Union (Maybe SymInteger)
          actual
            .@?= mrgIf
              (a .== 0)
              (return $ Just a)
              ( mrgIf
                  (b .== 0)
                  (return $ Just b)
                  (mrgIf (c .== 0) (return $ Just c) (return Nothing))
              )
      ]
  