Helper monad for cached computations on data that may change in small delta

CalCache.hs

{-# LANGUAGE FlexibleContexts, FlexibleInstances, GADTs, TypeFamilies #-}

module CalCache
  ( Ref(..)
  , Calc, NoCache, Cache
  , get, hasEq
  , toCache, updateCache
  , calcCacheResult
  ) where

import Control.Monad (liftM)
import Data.Set (Set)
import qualified Data.Set as Set

class (Monad (RefMonad r), Ord (RefIdent r)) => Ref r where
  type RefMonad r :: * -> *
  type RefIdent r :: *
  ident :: r a -> RefIdent r
  getRef :: r a -> RefMonad r a

data NoCache (r :: * -> *) a = NoCache

data MaybeEq a where
  NothingEq :: MaybeEq a
  JustEq :: Eq a => MaybeEq a

data Calc t r a = Calc
  { cache :: t r a
  , maybeEq :: MaybeEq a
  , action :: Action t r a
  }

hasEq :: Eq a => Calc t r a -> Calc t r a
hasEq c = c { maybeEq = JustEq }

data Action t r a where
  Return :: a -> Action t r a
  -- For Bind, the Calc.cache is for the second action
  Bind :: Calc t r a -> (a -> Calc NoCache r b) -> Action t r b
  Get :: r a -> Action t r a

instance Ref r => Monad (Calc NoCache r) where
  return = Calc NoCache NothingEq . Return
  (>>=) x = Calc NoCache NothingEq . Bind x

get :: Ref r => r a -> Calc NoCache r a
get = Calc NoCache NothingEq . Get

data Cache r a = Cache
  { sources :: Set (RefIdent r)
  , result :: a
  }

calcCacheResult :: Calc Cache r a -> a
calcCacheResult = result . cache

toCache :: Ref r => Calc NoCache r a -> RefMonad r (Calc Cache r a)
toCache c =
  case action c of
  Return x -> retRes Set.empty x $ Return x
  Get ref -> do
    r <- getRef ref
    retRes (Set.singleton (ident ref)) r $ Get ref
  Bind x y -> do
    newX <- toCache x
    newY <- toCache . y $ calcCacheResult newX
    let yCache = cache newY
    retRes (sources yCache) (result yCache) $ Bind newX y
  where
    retRes set res = return . Calc (Cache set res) (maybeEq c)

updateCacheH :: Ref r => Set (RefIdent r) -> Calc Cache r a -> RefMonad r (Calc Cache r a, Bool)
updateCacheH dirty c =
  case action c of
  Return x -> return $ mkRes False x
  Get ref
    | isDirty -> liftM (mkRes True) $ getRef ref
    | otherwise -> retCached
  Bind x y -> do
    (newX, xChanged) <- updateCacheH dirty x
    if xChanged || isDirty
      then
        liftM (mkRes True . calcCacheResult) .
        toCache . y $ calcCacheResult newX
      else
        retCached
  where
    retCached = return . mkRes False . result $ cache c
    mkRes mayChange res =
      ( Calc (Cache (sources (cache c)) res) (maybeEq c) (action c)
      , mayChange &&
        case maybeEq c of
        NothingEq -> True
        JustEq -> res /= result (cache c)
      )
    isDirty = not . Set.null . Set.intersection dirty . sources $ cache c

updateCache :: Ref r => Set (RefIdent r) -> Calc Cache r a -> RefMonad r (Calc Cache r a)
updateCache dirty = liftM fst . updateCacheH dirty