    module Futhark.AD.DerivativesTests (tests) where

    import Data.Map qualified as M
    import Data.Text qualified as T
    import Futhark.AD.Derivatives
    import Futhark.Analysis.PrimExp
    import Futhark.IR.Syntax.Core (nameFromText)
    import Futhark.Util.Pretty (prettyString)
    import Test.Tasty
    import Test.Tasty.HUnit

    tests :: TestTree
    tests :: TestTree
    tests =
      TestName -> [TestTree] -> TestTree
    testGroup
        TestName
    "Futhark.AD.DerivativesTests"
        [ TestName -> [TestTree] -> TestTree
    testGroup TestName
    "Primitive functions" ([TestTree] -> TestTree) -> [TestTree] -> TestTree
    forall a b. (a -> b) -> a -> b
    $
            ((Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
     -> TestTree)
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    -> [TestTree]
    forall a b. (a -> b) -> [a] -> [b]
    map (Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
    -> TestTree
    forall {c}. (Text, ([PrimType], PrimType, c)) -> TestTree
    primFunTest ([(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
     -> [TestTree])
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    -> [TestTree]
    forall a b. (a -> b) -> a -> b
    $
              ((Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
     -> Bool)
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool -> Bool
    not (Bool -> Bool)
    -> ((Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
        -> Bool)
    -> (Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Text -> [Text] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [Text]
    missing_primfuns) (Text -> Bool)
    -> ((Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
        -> Text)
    -> (Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
    -> Text
    forall a b. (a, b) -> a
    fst) ([(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
     -> [(Text,
          ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))])
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    forall a b. (a -> b) -> a -> b
    $
                Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
    -> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
    forall k a. Map k a -> [(k, a)]
    M.toList Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
    primFuns,
          TestName -> [TestTree] -> TestTree
    testGroup TestName
    "BinOps" ([TestTree] -> TestTree) -> [TestTree] -> TestTree
    forall a b. (a -> b) -> a -> b
    $ (BinOp -> TestTree) -> [BinOp] -> [TestTree]
    forall a b. (a -> b) -> [a] -> [b]
    map BinOp -> TestTree
    binOpTest [BinOp]
    allBinOps,
          TestName -> [TestTree] -> TestTree
    testGroup TestName
    "UnOps" ([TestTree] -> TestTree) -> [TestTree] -> TestTree
    forall a b. (a -> b) -> a -> b
    $ (UnOp -> TestTree) -> [UnOp] -> [TestTree]
    forall a b. (a -> b) -> [a] -> [b]
    map UnOp -> TestTree
    unOpTest [UnOp]
    allUnOps
        ]
      where
        blank :: PrimType -> PrimExp v
    blank = PrimValue -> PrimExp v
    forall v. PrimValue -> PrimExp v
    ValueExp (PrimValue -> PrimExp v)
    -> (PrimType -> PrimValue) -> PrimType -> PrimExp v
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PrimType -> PrimValue
    blankPrimValue

        primFunTest :: (Text, ([PrimType], PrimType, c)) -> TestTree
    primFunTest (Text
    f, ([PrimType]
    ts, PrimType
    ret, c
    _)) =
          TestName -> Assertion -> TestTree
    testCase (Text -> TestName
    T.unpack Text
    f) (Assertion -> TestTree) -> Assertion -> TestTree
    forall a b. (a -> b) -> a -> b
    $
            case Name -> [PrimExp VName] -> Maybe [PrimExp VName]
    pdBuiltin (Text -> Name
    nameFromText Text
    f) ((PrimType -> PrimExp VName) -> [PrimType] -> [PrimExp VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PrimType -> PrimExp VName
    forall {v}. PrimType -> PrimExp v
    blank [PrimType]
    ts) of
              Maybe [PrimExp VName]
    Nothing -> TestName -> Assertion
    forall a. HasCallStack => TestName -> IO a
    assertFailure TestName
    "pdBuiltin gives Nothing"
              Just [PrimExp VName]
    v -> (PrimExp VName -> PrimType) -> [PrimExp VName] -> [PrimType]
    forall a b. (a -> b) -> [a] -> [b]
    map PrimExp VName -> PrimType
    forall v. PrimExp v -> PrimType
    primExpType [PrimExp VName]
    v [PrimType] -> [PrimType] -> Assertion
    forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
    @?= Int -> PrimType -> [PrimType]
    forall a. Int -> a -> [a]
    replicate ([PrimType] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [PrimType]
    ts) PrimType
    ret

        -- We know we have no derivatives for these... and they are not
        -- coming any time soon.
        missing_primfuns :: [Text]
    missing_primfuns =
          [ Text
    "gamma16",
            Text
    "gamma32",
            Text
    "gamma64",
            Text
    "lgamma16",
            Text
    "lgamma32",
            Text
    "lgamma64"
          ]

        binOpTest :: BinOp -> TestTree
    binOpTest BinOp
    bop =
          TestName -> Assertion -> TestTree
    testCase (BinOp -> TestName
    forall a. Pretty a => a -> TestName
    prettyString BinOp
    bop) (Assertion -> TestTree) -> Assertion -> TestTree
    forall a b. (a -> b) -> a -> b
    $
            let t :: PrimType
    t = BinOp -> PrimType
    binOpType BinOp
    bop
                (PrimExp VName
    dx, PrimExp VName
    dy) = BinOp
    -> PrimExp VName -> PrimExp VName -> (PrimExp VName, PrimExp VName)
    pdBinOp BinOp
    bop (PrimType -> PrimExp VName
    forall {v}. PrimType -> PrimExp v
    blank PrimType
    t) (PrimType -> PrimExp VName
    forall {v}. PrimType -> PrimExp v
    blank PrimType
    t)
             in (PrimExp VName -> PrimType
    forall v. PrimExp v -> PrimType
    primExpType PrimExp VName
    dx, PrimExp VName -> PrimType
    forall v. PrimExp v -> PrimType
    primExpType PrimExp VName
    dy) (PrimType, PrimType) -> (PrimType, PrimType) -> Assertion
    forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
    @?= (PrimType
    t, PrimType
    t)

        unOpTest :: UnOp -> TestTree
    unOpTest UnOp
    bop =
          TestName -> Assertion -> TestTree
    testCase (UnOp -> TestName
    forall a. Pretty a => a -> TestName
    prettyString UnOp
    bop) (Assertion -> TestTree) -> Assertion -> TestTree
    forall a b. (a -> b) -> a -> b
    $
            let t :: PrimType
    t = UnOp -> PrimType
    unOpType UnOp
    bop
             in PrimExp VName -> PrimType
    forall v. PrimExp v -> PrimType
    primExpType (UnOp -> PrimExp VName -> PrimExp VName
    pdUnOp UnOp
    bop (PrimExp VName -> PrimExp VName) -> PrimExp VName -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $ PrimType -> PrimExp VName
    forall {v}. PrimType -> PrimExp v
    blank PrimType
    t) PrimType -> PrimType -> Assertion
    forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
    @?= PrimType
    t