{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, Arrows #-}

module Control.CacheArrow ( CacheArrow
                          , runCache'
                          , runCache
                          ) where

import Prelude hiding ((.), id)
import Control.Arrow
import Control.Arrow.Transformer
import Control.Category


data CacheArrow a b c = CacheArrow (Maybe (b, c)) (a b c)

instance Arrow a => Category (CacheArrow a) where
  (CacheArrow _ g) . (CacheArrow _ f) = CacheArrow Nothing (g . f)

  id = arr id

instance Arrow a => Arrow (CacheArrow a) where
  arr = lift . arr
  first (CacheArrow _ f) = lift . first $ f
  second (CacheArrow _ f) = lift . second $ f

instance ArrowChoice a => ArrowChoice (CacheArrow a) where
  left (CacheArrow _ f) = lift . left $ f
  right (CacheArrow _ f) = lift . right $ f

instance Arrow a => ArrowTransformer CacheArrow a where
  lift = CacheArrow Nothing


runCacheAlways :: Arrow a => CacheArrow a b c -> a b (c, CacheArrow a b c)
runCacheAlways (CacheArrow _ f) = proc b -> do
  c <- f -< b
  returnA -< (c, CacheArrow (Just (b, c)) f)

runCache' :: (ArrowChoice a, Eq b) => CacheArrow a b c -> a b (c, Bool, CacheArrow a b c)
runCache' a@(CacheArrow cache _) = proc b ->
  case cache of
    (Just (b', c')) | b == b' -> returnA -< (c', False, a)
    _ -> do
      (c, cache') <- runCacheAlways a -< b
      returnA -< (c, True, cache')

runCache :: (ArrowChoice a, Eq b) => CacheArrow a b c -> a b (c, CacheArrow a b c)
runCache a = runCache' a >>^ (\(c, _, cache) -> (c, cache))