Specialised requests

Haxl/Core/Monad.hs

@@ -31,6 +31,7 @@ module Haxl.Core.Monad (
 
  -- * Data fetching and caching
  dataFetch, uncachedRequest,
+ dataFetchSpecialised,
  cacheRequest, cacheResult, cachedComputation,
  dumpCacheAsHaskell,
 
@@ -79,7 +80,7 @@ import Control.Arrow (left)
 import Control.Exception (bracket_)
 import Debug.Trace (traceEventIO)
 #endif
-
+import Unsafe.Coerce
 -- -----------------------------------------------------------------------------
 -- The environment
 
@@ -409,6 +410,56 @@ dataFetch req = GenHaxl $ \env ref -> do
  Cached (Left ex) -> return (Throw ex)
  Cached (Right a) -> return (Done a)
 
+-- | Sometimes, one request is a special case of a more general
+-- request. In case the general case has already been requested, we
+-- can avoid scheduling the special request, and determine its result
+-- from the result of the more general request.
+--
+-- An example would be determining the number of some xs (the special
+-- case) and retrieving the xs (the general case).
+dataFetchSpecialised :: forall r a u b . (DataSource u r, Request r a, Request r b)
+ => (r a -> r b) -- ^ given a request, determine a more general request
+ -> (b -> a) -- ^ how to ge the special result from the more genereal result
+ -> r a -- ^ (special) request
+ -> GenHaxl u a
+dataFetchSpecialised generalise specialise req = GenHaxl $ \ env ref -> do
+ -- first, check if a more general request has already been
+ -- performed. We don't want to use 'cached' for this, since we
+ -- don't want to add it if it's not there.
+ cache <- readIORef (cacheRef env)
+ case DataCache.lookup (generalise req) cache of
+ Nothing -> do
+ allRequests <- readIORef ref
+ let test (BlockedFetch r _) = unsafeCoerce r == generalise req
+ req' = find test
+ (requestsOfType (generalise req) allRequests)
+ case req' of
+ Just (BlockedFetch r rvar) ->
+ let r' = unsafeCoerce r :: r b
+ rvar' = unsafeCoerce rvar
+ in return $ Blocked (Cont (continueFetch' specialise r' rvar'))
+ Nothing ->
+ -- the more general request is neither in the cache, nor in
+ -- the RequestStore, just perform the request
+ (unHaxl $ dataFetch req) env ref
+ Just rvar -> do
+ mb <- tryReadResult rvar
+ case mb of
+ Just (Right r) -> return (Done (specialise r))
+ Just (Left ex) -> return (Throw ex)
+ Nothing -> return $
+ Blocked (Cont (continueFetch' specialise (generalise req) rvar))
+
+continueFetch'
+ :: (DataSource u r, Request r b, Show b)
+ => (b -> a) -> r b -> ResultVar b -> GenHaxl u a
+continueFetch' f req rvar = GenHaxl $ \_env _ref -> do
+ m <- tryReadResult rvar
+ case m of
+ Nothing -> raise . DataSourceError $
+ textShow req <> " did not set contents of result var"
+ Just r -> done (f <$> r)
+
 -- | A data request that is not cached. This is not what you want for
 -- normal read requests, because then multiple identical requests may
 -- return different results, and this invalidates some of the

Haxl/Core/RequestStore.hs

@@ -8,7 +8,7 @@
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-
+{-# LANGUAGE ConstraintKinds #-}
 -- | Bucketing requests by 'DataSource'.
 --
 -- When a request is issued by the client via 'dataFetch', it is placed
@@ -20,7 +20,8 @@
 -- users should not need to import it.
 module Haxl.Core.RequestStore (
  BlockedFetches(..), RequestStore,
- noRequests, addRequest, contents
+ noRequests, addRequest, contents,
+ requestsOfType
  ) where
 
 import Haxl.Core.Types
@@ -70,3 +71,12 @@ addRequest bf (RequestStore m) =
 -- | Retrieves the whole contents of the 'RequestStore'.
 contents :: RequestStore u -> [BlockedFetches u]
 contents (RequestStore m) = Map.elems m
+
+-- | Retrieves requests in the 'RequestStore' that have the same type
+-- as a given request.
+requestsOfType :: forall r a u . (DataSource u r, Request r a) => r a -> RequestStore u -> [BlockedFetch r]
+requestsOfType _ (RequestStore rs) =
+ let ty = typeOf1 (undefined :: r a)
+ in case Map.lookup ty rs of
+ Just (BlockedFetches result) -> map unsafeCoerce result
+ Nothing -> []

tests/ExampleDataSource.hs

@@ -1,11 +1,11 @@
-{-# LANGUAGE DeriveDataTypeable #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE StandaloneDeriving #-}
-{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module ExampleDataSource (
  -- * initialise the state
@@ -15,17 +15,19 @@ module ExampleDataSource (
  Id(..), ExampleReq(..),
  countAardvarks,
  listWombats,
+ countWombats
  ) where
 
-import Haxl.Prelude
-import Prelude ()
+import  Haxl.Prelude
+import  Prelude  ()
 
-import Haxl.Core
+import Haxl.Core
+import Haxl.Core.Monad (dataFetchSpecialised)
 
-import Data.Typeable
-import Data.Hashable
-import Control.Concurrent
-import System.IO
+import  Control.Concurrent
+import  Data.Hashable
+import  Data.Typeable
+import  System.IO
 
 -- Here is an example minimal data source. Our data source will have
 -- two requests:
@@ -45,6 +47,7 @@ instance Show Id where
 data ExampleReq a where
  CountAardvarks :: String -> ExampleReq Int
  ListWombats :: Id -> ExampleReq [Id]
+ CountWombats :: Id -> ExampleReq Int
  deriving Typeable -- requests must be Typeable
 
 -- The request type (ExampleReq) is parameterized by the result type of
@@ -70,6 +73,7 @@ instance Show1 ExampleReq where show1 = show
 instance Hashable (ExampleReq a) where
  hashWithSalt s (CountAardvarks a) = hashWithSalt s (0::Int,a)
  hashWithSalt s (ListWombats a) = hashWithSalt s (1::Int,a)
+ hashWithSalt s (CountWombats a) = hashWithSalt s (2::Int,a)
 
 instance StateKey ExampleReq where
  data State ExampleReq = ExampleState {
@@ -145,7 +149,10 @@ fetch1 (BlockedFetch (ListWombats a) r) =
  if a > 999999
  then putFailure r $ FetchError "too large"
  else putSuccess r $ take (fromIntegral a) [1..]
-
+fetch1 (BlockedFetch (CountWombats a) r) =
+ if a > 999999
+ then putFailure r $ FetchError "too large"
+ else putSuccess r $ fromIntegral a
 
 -- Normally a data source will provide some convenient wrappers for
 -- its requests:
@@ -155,3 +162,11 @@ countAardvarks str = dataFetch (CountAardvarks str)
 
 listWombats :: Id -> GenHaxl () [Id]
 listWombats id = dataFetch (ListWombats id)
+
+countWombats :: Id -> GenHaxl () Int
+countWombats id =
+ let f :: ExampleReq Int -> ExampleReq [Id]
+ f (CountWombats a) = ListWombats a
+ g :: [Id] -> Int
+ g = length
+ in dataFetchSpecialised f g (CountWombats id)

tests/TestExampleDataSource.hs

@@ -33,7 +33,8 @@ tests = TestList [
  TestLabel "cachedComputationTest" cachedComputationTest,
  TestLabel "memoTest" memoTest,
  TestLabel "dataSourceExceptionTest" dataSourceExceptionTest,
- TestLabel "dumpCacheAsHaskell" dumpCacheTest]
+ TestLabel "dumpCacheAsHaskell" dumpCacheTest,
+ TestLabel "specialisedTest" specialisedTest]
 
 -- Let's test ExampleDataSource.
 
@@ -175,3 +176,12 @@ dumpCacheTest = TestCase $ do
  str <- runHaxl env dumpCacheAsHaskell
  loadcache <- readFile "sigma/haxl/core/tests/LoadCache.txt"
  assertEqual "dumpCacheAsHaskell" str loadcache
+
+specialisedTest = TestCase $ do
+ env <- testEnv
+ let env' = env { flags = (flags env){trace = 3} }
+ wombats <- runHaxl env' $ listWombats 5
+ nWombats <- runHaxl env' $ countWombats 5
+ assertEqual "nWombats" (length wombats) nWombats
+ stats <- readIORef (statsRef env)
+ assertEqual "fetches" 1 (numFetches stats)