From 762d1100414fd5b365b496669eeaef6846609f1a Mon Sep 17 00:00:00 2001
From: Simon Marlow <smarlow@fb.com>
Date: Tue, 19 Sep 2017 07:33:38 -0700
Subject: [PATCH] Add preCacheComputation

Summary:
Restricted version of cachedComputation, useful for abusing
memoization for mocking functions.

Reviewed By: niteria

Differential Revision: D5861480

fbshipit-source-id: c3978bb675e260591bb99b7dd2219775a9422907
---
 Haxl/Core.hs       |  3 ++-
 Haxl/Core/Monad.hs | 30 +++++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/Haxl/Core.hs b/Haxl/Core.hs
index 3399e3f..8ab50fc 100644
--- a/Haxl/Core.hs
+++ b/Haxl/Core.hs
@@ -26,7 +26,8 @@ module Haxl.Core (
 
   -- ** Data fetching and caching
   dataFetch, uncachedRequest,
-  cacheRequest, cacheResult, cacheResultWithShow, cachedComputation,
+  cacheRequest, cacheResult, cacheResultWithShow,
+  cachedComputation, preCacheComputation,
   dumpCacheAsHaskell,
 
   -- ** Memoization
diff --git a/Haxl/Core/Monad.hs b/Haxl/Core/Monad.hs
index c06abc4..b60e585 100644
--- a/Haxl/Core/Monad.hs
+++ b/Haxl/Core/Monad.hs
@@ -40,7 +40,7 @@ module Haxl.Core.Monad (
 
     -- * Data fetching and caching
     ShowReq, dataFetch, dataFetchWithShow, uncachedRequest, cacheRequest,
-    cacheResult, cacheResultWithShow, cachedComputation,
+    cacheResult, cacheResultWithShow, cachedComputation, preCacheComputation,
     dumpCacheAsHaskell, dumpCacheAsHaskellFn,
 
     -- * Memoization Machinery
@@ -988,6 +988,34 @@ cachedComputation req haxl = do
 done :: Either SomeException a -> IO (Result u a)
 done = return . either Throw Done
 
+-- | Like 'cachedComputation', but fails if the cache is already
+-- populated.
+--
+-- Memoization can be (ab)used to "mock" a cached computation, by
+-- pre-populating the cache with an alternative implementation. In
+-- that case we don't want the operation to populate the cache to
+-- silently succeed if the cache is already populated.
+--
+preCacheComputation
+  :: forall req u a.
+     ( Eq (req a)
+     , Hashable (req a)
+     , Typeable (req a))
+  => req a -> GenHaxl u a -> GenHaxl u a
+preCacheComputation req haxl = do
+  env <- env id
+  cache <- unsafeLiftIO $ readIORef (memoRef env)
+  unsafeLiftIO $ ifProfiling (flags env) $
+    modifyIORef' (profRef env) (incrementMemoHitCounterFor (profLabel env))
+  case DataCache.lookup req cache of
+    Nothing -> do
+      memoVar <- newMemoWith haxl
+      unsafeLiftIO $ writeIORef (memoRef env) $!
+        DataCache.insertNotShowable req memoVar cache
+      runMemo memoVar
+    Just _ -> throw $ InvalidParameter $
+      "preCacheComputation: key is already cached"
+
 -- -----------------------------------------------------------------------------
 
 -- | Dump the contents of the cache as Haskell code that, when