Implementing a high-order Cache effect - Polysemy

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

import Data.Hashable (Hashable)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import Polysemy
import Polysemy.State
-- import Polysemy.AtomicState -- TODO: thread safety

data Cache k v m a where
  CacheGet
    :: (k -> m v) -- ^ How to retrieve a value if it doesn't exist in the cache
    -> (v -> m v) -- ^ How to return a value if it does exist in the cache (useful to e.g. add logging, otherwise just use 'pure'
    -> k -- ^ The key to retrieve
    -> Cache k v m v -- ^ Returns the value

makeSem ''Cache

runCacheAsHashMap
  :: forall k v r a
   . (Eq k, Hashable k)
  => Sem (Cache k v ': r) a -> Sem r a
runCacheAsHashMap = evalState (HM.empty @k @v) . reinterpretH
  (\case
    CacheGet fetch returnFromCache k -> do
      cache <- get
      case HM.lookup k cache of
        Just v -> do
          returnFromCache' <- bindT returnFromCache -- f v -> Sem (Cache k v ': State (HashMap k v) ': r) (f v)
          pureT v >>= returnFromCache'
        Nothing -> do
          fetch' <- bindT fetch                     -- f k -> Sem (Cache k v ': State (HashMap k v) ': r) (f v)
          fv <- pureT k >>= fetch'
          let fcache = (\v' -> HM.insert k v' cache) <$> fv
          put' <- bindT put
          put' fcache
          pure fv
  )

src/Cache.hs:61:25-27: error:
    • Ambiguous use of effect 'State'
      Possible fix:
        add (Member (State (HashMap k v)) r0) to the context of
          the type signature
      If you already have the constraint you want, instead
        add a type application to specify
          'HashMap k v' directly, or activate polysemy-plugin which
            can usually infer the type correctly.
    • In the first argument of ‘bindT’, namely ‘put’
      In a stmt of a 'do' block: put' <- bindT put
      In the expression:
        do fetch' <- bindT fetch
           fv <- pureT k >>= fetch'
           let fcache = (\ v' -> ...) <$> fv
           put' <- bindT put
           ....
   |
61 |           put' <- bindT put
   |                         ^^^

data Cache k v m a where
  Cache :: k -> m v -> Cache k v m v

runCache
  :: (Ord k, Member (State (M.Map k v)) r)
  => Sem (Cache k v : r) a -> Sem r a
runCache = interpretH \(Cache k mv) ->
  gets (M.!? k) >>= \case
    Just v  -> pureT v
    Nothing -> do
      fv <- join $ raise . runCache <$> runT mv
      (`inspect` fv) <$> getInspectorT >>=
        maybe (pure ()) (modify . M.insert k)
      pure fv

runCacheInHashMap
  :: (Eq k, Hashable k, Member (State (HashMap k v)) r)
  => Sem (Cache k v : r) a -> Sem r a
runCacheInHashMap = interpretH $ \(Cache k mv) ->
  gets (HM.lookup k) >>= \case
    Just v  -> pureT v
    Nothing -> do
      fv <- runT mv >>= raise . runCacheInHashMap
      (`inspect` fv) <$> getInspectorT >>=
        maybe (pure ()) (modify . HM.insert k)
      pure fv

runCacheAsHashMap
  :: (Eq k, Hashable k)
  => Sem (Cache k v ': r) a -> Sem r a
runCacheAsHashMap = evalState HM.empty . runCacheInHashMap . raiseUnder