content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TypeFamilies #-}
  module WriteTests (tests) where
  
  import Test.HUnit
  
  import Control.Arrow
  import Control.Concurrent
  import Data.Either
  import Data.Foldable
  import Data.Hashable
  import Data.IORef
  import qualified Data.Text as Text
  
  import Haxl.Core.Monad (mapWrites, mapWriteTree, flattenWT, WriteTree)
  import Haxl.Core
  import Haxl.Prelude as Haxl
  
  -- A fake data source
  data SimpleDataSource a where
    GetNumber :: SimpleDataSource Int
  
  deriving instance Eq (SimpleDataSource a)
  deriving instance Show (SimpleDataSource a)
  instance ShowP SimpleDataSource where showp = show
  
  instance Hashable (SimpleDataSource a) where
    hashWithSalt s GetNumber = hashWithSalt s (0 :: Int)
  
  instance StateKey SimpleDataSource where
    data State SimpleDataSource = DSState
  
  instance DataSourceName SimpleDataSource where
    dataSourceName _ = "SimpleDataSource"
  
  instance DataSource u SimpleDataSource where
    fetch _st _flags _usr  = SyncFetch $ Haxl.mapM_ fetch1
      where
      fetch1 :: BlockedFetch SimpleDataSource -> IO ()
      fetch1 (BlockedFetch GetNumber m) =
        threadDelay 1000 >> putSuccess m 37
  
  newtype SimpleWrite = SimpleWrite Text
    deriving (Eq, Show, Ord)
  
  assertEqualIgnoreOrder ::
    (Eq a, Show a, Ord a) => String -> [a] -> [a] -> Assertion
  assertEqualIgnoreOrder msg lhs rhs =
    assertEqual msg (sort lhs) (sort rhs)
  
  doInnerWrite :: GenHaxl u (WriteTree SimpleWrite) Int
  doInnerWrite = do
    tellWrite $ SimpleWrite "inner"
    return 0
  
  doOuterWrite :: GenHaxl u (WriteTree SimpleWrite) Int
  doOuterWrite = do
    tellWrite $ SimpleWrite "outer1"
  
    doWriteMemo <- newMemoWith doInnerWrite
    let doMemoizedWrite = runMemo doWriteMemo
    _ <- doMemoizedWrite
    _ <- doMemoizedWrite
  
    tellWrite $ SimpleWrite "outer2"
  
    return 1
  
  doNonMemoWrites :: GenHaxl u (WriteTree SimpleWrite) Int
  doNonMemoWrites = do
    tellWrite $ SimpleWrite "inner"
    tellWriteNoMemo $ SimpleWrite "inner not memo"
    return 0
  
  runHaxlWithWriteList :: Env u (WriteTree w) -> GenHaxl u (WriteTree w) a -> IO (a, [w])
  runHaxlWithWriteList env haxl = second flattenWT <$> runHaxlWithWrites env haxl
  
  writeSoundness :: Test
  writeSoundness = TestCase $ do
    let numReps = 4
  
    -- do writes without memoization
    env1 <- emptyEnv ()
    (allRes, allWrites) <- runHaxlWithWriteList env1 $
      Haxl.sequence (replicate numReps doInnerWrite)
  
    assertBool "Write Soundness 1" $
      allWrites == replicate numReps (SimpleWrite "inner")
    assertBool "Write Soundness 2" $ allRes == replicate numReps 0
  
    -- do writes with memoization
    env2 <- emptyEnv ()
  
    (memoRes, memoWrites) <- runHaxlWithWriteList env2 $ do
      doWriteMemo <- newMemoWith doInnerWrite
      let memoizedWrite = runMemo doWriteMemo
  
      Haxl.sequence (replicate numReps memoizedWrite)
  
    assertBool "Write Soundness 3" $
      memoWrites == replicate numReps (SimpleWrite "inner")
    assertBool "Write Soundness 4" $ memoRes == replicate numReps 0
  
    -- do writes with interleaved memo
    env3 <- emptyEnv ()
  
    (ilRes, ilWrites) <- runHaxlWithWriteList env3 $ do
      doWriteMemo <- newMemoWith doInnerWrite
      let memoizedWrite = runMemo doWriteMemo
  
      Haxl.sequence $ replicate numReps (doInnerWrite *> memoizedWrite)
  
    assertBool "Write Soundness 5" $
      ilWrites == replicate (2*numReps) (SimpleWrite "inner")
    assertBool "Write Soundness 6" $ ilRes == replicate numReps 0
  
    -- do writes with nested memo
    env4 <- emptyEnv ()
  
    (nestRes, nestWrites) <- runHaxlWithWriteList env4 $ do
      doWriteMemo' <- newMemoWith doOuterWrite
      let memoizedWrite' = runMemo doWriteMemo'
  
      Haxl.sequence (replicate numReps memoizedWrite')
  
    let expWrites =
          [ SimpleWrite "outer1"
          , SimpleWrite "inner"
          , SimpleWrite "inner"
          , SimpleWrite "outer2"
          ]
    assertBool "Write Soundness 7" $
      nestWrites == fold (replicate numReps expWrites)
    assertBool "Write Soundness 8" $ nestRes == replicate numReps 1
  
    -- do both kinds of writes without memoization
    env5 <- emptyEnv ()
    (allRes, allWrites) <- runHaxlWithWriteList env5 $
      Haxl.sequence (replicate numReps doNonMemoWrites)
  
    assertBool "Write Soundness 9" $
      allWrites == replicate numReps (SimpleWrite "inner") ++
        replicate numReps (SimpleWrite "inner not memo")
    assertBool "Write Soundness 10" $ allRes == replicate numReps 0
  
    -- do both kinds of writes with memoization
    env6 <- emptyEnv ()
  
    (memoRes, memoWrites) <- runHaxlWithWriteList env6 $ do
      doWriteMemo <- newMemoWith doNonMemoWrites
      let memoizedWrite = runMemo doWriteMemo
  
      Haxl.sequence (replicate numReps memoizedWrite)
  
    -- "inner not memo" only appears once in this test
    assertBool "Write Soundness 11" $
      memoWrites == replicate numReps (SimpleWrite "inner") ++
        [SimpleWrite "inner not memo"]
    assertBool "Write Soundness 12" $ memoRes == replicate numReps 0
  
  writeLogsCorrectnessTest :: Test
  writeLogsCorrectnessTest = TestLabel "writeLogs_correctness" $ TestCase $ do
    e <- emptyEnv ()
    (_ , wrts) <- runHaxlWithWriteList e doNonMemoWrites
    assertEqualIgnoreOrder "Expected writes" [SimpleWrite "inner",
      SimpleWrite "inner not memo"] wrts
    wrtsNoMemo <- readIORef $ writeLogsRefNoMemo e
    wrtsMemo <- readIORef $ writeLogsRef e
    assertEqualIgnoreOrder "WriteTree not empty" [] $ flattenWT wrtsNoMemo
    assertEqualIgnoreOrder "WriteTree not empty" [] $ flattenWT wrtsMemo
  
  mapWritesTest :: Test
  mapWritesTest = TestLabel "mapWrites" $ TestCase $ do
    let funcSingle (SimpleWrite s) = SimpleWrite $ Text.toUpper s
        func = mapWriteTree funcSingle
    env0 <- emptyEnv ()
    (res0, wrts0) <- runHaxlWithWriteList env0 $ mapWrites func doNonMemoWrites
    assertEqual "Expected computation result" 0 res0
    assertEqualIgnoreOrder "Writes correctly transformed" [SimpleWrite "INNER",
      SimpleWrite "INNER NOT MEMO"] wrts0
  
    -- Writes should behave the same inside and outside mapWrites
    env1 <- emptyEnv ()
    (res1, wrts1) <- runHaxlWithWriteList env1 $ do
      outer <- doOuterWrite
      outerMapped <- mapWrites func doOuterWrite
      return $ outer == outerMapped
    assertBool "Results are identical" res1
    assertEqualIgnoreOrder
      "Writes correctly transformed, non-transformed writes preserved"
      [ SimpleWrite "outer1", SimpleWrite "inner"
      , SimpleWrite "inner", SimpleWrite "outer2"
      , SimpleWrite "OUTER1", SimpleWrite "INNER"
      , SimpleWrite "INNER", SimpleWrite "OUTER2"
      ]
      wrts1
  
    -- Memoization behaviour should be unaffected
    env2 <- emptyEnv ()
    (_res2, wrts2) <- runHaxlWithWriteList env2 $ do
      writeMemo <- newMemoWith doNonMemoWrites
      let doWriteMemo = runMemo writeMemo
      _ <- mapWrites func doWriteMemo
      _ <- doWriteMemo
      return ()
    -- "inner not memo" should appear only once
    assertEqualIgnoreOrder
      "Write correctly transformed under memoization"
      [ SimpleWrite "INNER"
      , SimpleWrite "inner"
      , SimpleWrite "INNER NOT MEMO"
      ]
      wrts2
  
    -- Same as previous, but the non-mapped computation is run first
    env3 <- emptyEnv ()
    (_res3, wrts3) <- runHaxlWithWriteList env3 $ do
      writeMemo <- newMemoWith doNonMemoWrites
      let doWriteMemo = runMemo writeMemo
      _ <- doWriteMemo
      _ <- mapWrites func doWriteMemo
      return ()
    -- "inner not memo" should appear only once
    assertEqualIgnoreOrder
      "Flipped: Write correctly transformed under memoization"
      [ SimpleWrite "inner"
      , SimpleWrite "INNER"
      , SimpleWrite "inner not memo"
      ]
      wrts3
  
    -- inner computation performs no writes
    env4 <- emptyEnv ()
    (res4, wrts4) <- runHaxlWithWriteList env4 $
      mapWrites func (return (0 :: Int))
    assertEqual "No Writes: Expected computation result" 0 res4
    assertEqualIgnoreOrder "No writes" [] wrts4
  
    -- inner computation throws an exception
    env5 <- emptyEnv ()
    (res5, wrts5) <- runHaxlWithWriteList env5 $ mapWrites func $ try $ do
      _ <- doNonMemoWrites
      _ <- throw (NotFound "exception")
      return 0
    assertBool "Throw: Expected Computation Result" $ isLeft
      (res5 :: Either HaxlException Int)
    assertEqualIgnoreOrder
      "Datasource writes correctly transformed"
      [ SimpleWrite "INNER"
      , SimpleWrite "INNER NOT MEMO"
      ]
      wrts5
  
    -- inner computation calls a datasource
    env6 <- initEnv (stateSet DSState stateEmpty) ()
    (res6, wrts6) <- runHaxlWithWriteList env6 $ mapWrites func $ do
      _ <- doNonMemoWrites
      dataFetch GetNumber
  
    assertEqual "Datasource: Expected Computation Result" 37 res6
    assertEqualIgnoreOrder
      "Datasource writes correctly transformed"
      [ SimpleWrite "INNER"
      , SimpleWrite "INNER NOT MEMO"
      ]
      wrts6
  
    -- inner computation calls a datasource, flipped calls
    env7 <- initEnv (stateSet DSState stateEmpty) ()
    (res7, wrts7) <- runHaxlWithWriteList env7 $ mapWrites func $ do
      df <- dataFetch GetNumber
      _ <- doNonMemoWrites
      return df
  
    assertEqual "Flipped Datasource: Expected Computation Result" 37 res7
    assertEqualIgnoreOrder
      "Flipped: Datasource writes correctly transformed"
      [ SimpleWrite "INNER"
      , SimpleWrite "INNER NOT MEMO"
      ]
      wrts7
  
  tests = TestList
    [ TestLabel "Write Soundness" writeSoundness,
      TestLabel "writeLogs_correctness" writeLogsCorrectnessTest,
      TestLabel "mapWrites" mapWritesTest
    ]
  