Generalize Haxl.Core.Monad.mapWrites

Summary:
On my last diff I forgot to generalize `mapWrites` in Haxl core, leaving it working only with `WriteTree`s. Mapping functionality is needed if we want to use alternative data structures.

Luckily, the function can be trivially generalized by mapping the `Monoid` instead and exposing `mapWriteTree` to lift a function mapping elements to one mapping `WriteTree`s.

Reviewed By: josefs

Differential Revision: D39500863

fbshipit-source-id: c0c140544e135814575bc1500fec649fabc1fdec

Haxl/Core/Monad.hs

@@ -49,6 +49,7 @@ module Haxl.Core.Monad
  , appendWTs
  , mbModifyWLRef
  , mapWrites
+ , mapWriteTree
 
  -- * Cont
  , Cont(..)
@@ -912,10 +913,10 @@ nextCallId env = atomicModifyIORef' (callIdRef env) $ \x -> (x+1,x+1)
 -- Memoization behavior is unchanged, meaning if a memoized computation is run
 -- once inside @mapWrites@ and then once without, the writes from the second run
 -- will NOT be transformed.
-mapWrites :: (w -> w) -> GenHaxl u (WriteTree w) a -> GenHaxl u (WriteTree w) a
+mapWrites :: Monoid w => (w -> w) -> GenHaxl u w a -> GenHaxl u w a
 mapWrites f action = GenHaxl $ \curEnv -> do
- wlogs <- newIORef NilWrites
- wlogsNoMemo <- newIORef NilWrites
+ wlogs <- newIORef mempty
+ wlogsNoMemo <- newIORef mempty
  let
  !newEnv = curEnv { writeLogsRef = wlogs, writeLogsRefNoMemo = wlogsNoMemo }
  unHaxl (mapWritesImpl curEnv newEnv action) newEnv
@@ -924,9 +925,9 @@ mapWrites f action = GenHaxl $ \curEnv -> do
  let
  pushTransformedWrites = do
  wt <- readIORef $ writeLogsRef curEnv
- mbModifyWLRef (mapWriteTree f wt) (writeLogsRef oldEnv)
+ modifyIORef' (writeLogsRef oldEnv) (<> f wt)
  wtNoMemo <- readIORef $ writeLogsRefNoMemo curEnv
- mbModifyWLRef (mapWriteTree f wtNoMemo) (writeLogsRefNoMemo oldEnv)
+ modifyIORef' (writeLogsRefNoMemo oldEnv) (<> f wtNoMemo)
 
  r <- m curEnv
 

tests/WriteTests.hs

@@ -23,7 +23,7 @@ import Data.Hashable
 import Data.IORef
 import qualified Data.Text as Text
 
-import Haxl.Core.Monad (mapWrites, flattenWT, WriteTree)
+import Haxl.Core.Monad (mapWrites, mapWriteTree, flattenWT, WriteTree)
 import Haxl.Core
 import Haxl.Prelude as Haxl
 
@@ -182,7 +182,8 @@ writeLogsCorrectnessTest = TestLabel "writeLogs_correctness" $ TestCase $ do
 
 mapWritesTest :: Test
 mapWritesTest = TestLabel "mapWrites" $ TestCase $ do
- let func (SimpleWrite s) = SimpleWrite $ Text.toUpper s
+ 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