1 files changed, 132 insertions, 1 deletions
diff --git a/src/lib/IncState.hs b/src/lib/IncState.hs
index 1f676546..b0b9e4ec 100644
--- a/src/lib/IncState.hs
+++ b/src/lib/IncState.hs
@@ -9,11 +9,142 @@
 module IncState (
   IncState (..), MapEltUpdate (..), MapUpdate (..),
   Overwrite (..), TailUpdate (..), Unchanging (..), Overwritable (..),
-  mapUpdateMapWithKey, MonoidState (..)) where
+  mapUpdateMapWithKey, MonoidState (..), AllOrNothing (..), fmapIncMap,
+  IncM, IncVar, liftIncM, runIncM, IncFun, fmapIncVar, incZip2, incUnzip3,
+  incUnzip2, incZip3, liftMonoidStateIncM) where
+
+import Control.Monad.State.Strict
+import Data.IORef
 
 import Data.Aeson (ToJSON (..))
 import qualified Data.Map.Strict as M
 import GHC.Generics
+import Data.Maybe (fromJust)
+
+-- === incremental computation builder ===
+
+-- We use IO here for IORefs but we could use ST or something else instead
+type IncFun a b = a -> IO (b, Delta a -> IO (Delta b))
+type IncM = StateT (IO ()) IO
+type IncVar a = (a, IORef (Maybe (Delta a)))
+
+liftIncM :: IncVar a -> IncFun a b -> IncM (IncVar b)
+liftIncM (x, dxRef) f = do
+  (y, df) <- liftIO $ f x
+  dyRef <- liftIO $ newIORef Nothing
+  addIncAction do
+    Just dx <- liftIO $ readIORef dxRef
+    dy <- df dx
+    liftIO $ writeIORef dyRef (Just dy)
+  return (y, dyRef)
+
+-- like LiftIncM but you don't have to bother with the initial values
+liftMonoidStateIncM :: IncVar (MonoidState a) -> IO (a -> IO b) -> IncM (IncVar (MonoidState b))
+liftMonoidStateIncM v createIncFun = liftIncM v \(MonoidState xInit) -> do
+  incFun <- createIncFun
+  yInit <- incFun xInit
+  return (MonoidState yInit, incFun)
+
+runIncM :: (IncVar a -> IncM (IncVar b)) -> IncFun a b
+runIncM f = \x -> do
+  dxRef <- newIORef Nothing
+  ((y, dyRef), action) <- runStateT (f (x, dxRef)) (return ())
+  return (y, \dx -> do
+     writeIORef dxRef (Just dx)
+     action
+     fromJust <$> readIORef dyRef)
+
+fmapIncVar :: IncVar a -> (a -> b) -> (Delta a -> Delta b) -> IncM (IncVar b)
+fmapIncVar v f df = liftIncM v \x -> return (f x, \dx -> return $ df dx)
+
+fmapIncMap
+  :: forall k a b. Ord k
+  => IncVar (M.Map k a) -> (k -> IncVar a -> IncM (IncVar b)) -> IncM (IncVar (M.Map k b))
+fmapIncMap v f = liftIncM v \m -> do
+  initDfsAndResults <- flip M.traverseWithKey m \k x -> runIncM (f k) x
+  let initResults = (fst <$> initDfsAndResults) :: M.Map k b
+  let initDfs     = (snd <$> initDfsAndResults) :: M.Map k (Delta a -> IO (Delta b))
+  dfsRef <- newIORef initDfs
+  return (initResults, deltaComputation dfsRef)
+  where
+    deltaComputation
+      :: IORef (M.Map k (Delta a -> IO (Delta b)))
+      -> MapUpdate k a -> IO (MapUpdate k b)
+    deltaComputation dfs dxs = MapUpdate <$> do
+      flip M.traverseWithKey (mapUpdates dxs) \k -> \case
+       Create x -> do
+         (y, df) <- runIncM (f k) x
+         modifyIORef dfs (M.insert k df)
+         return $ Create y
+       Replace x -> do
+         (y, df) <- runIncM (f k) x
+         modifyIORef dfs (M.insert k df)
+         return $ Replace y
+       Update dx -> do
+         df <- fromJust <$> M.lookup k <$> readIORef dfs
+         Update <$> df dx
+       Delete -> do
+         modifyIORef dfs (M.delete k)
+         return Delete
+
+incUnzip2 :: IncVar (a, b) -> IncM (IncVar a, IncVar b)
+incUnzip2 v = do
+  x <- fmapIncVar v (\(x, _) -> x) (\(dx, _ ) -> dx)
+  y <- fmapIncVar v (\(_, y) -> y) (\(_ , dy) -> dy)
+  return (x, y)
+
+incUnzip3 :: IncVar (a, b, c) -> IncM (IncVar a, IncVar b, IncVar c)
+incUnzip3 v = do
+  x <- fmapIncVar v (\(x, _, _) -> x) (\(dx, _ , _ ) -> dx)
+  y <- fmapIncVar v (\(_, y, _) -> y) (\(_ , dy, _ ) -> dy)
+  z <- fmapIncVar v (\(_, _, z) -> z) (\(_ , _ , dz) -> dz)
+  return (x, y, z)
+
+zipIncVar :: IncVar a -> IncVar b -> IncM (IncVar (a, b))
+zipIncVar (x, dxRef) (y, dyRef) = do
+  let xy = (x, y)
+  dxyRef <- liftIO $ newIORef Nothing
+  addIncAction do
+    Just dx <- liftIO $ readIORef dxRef
+    Just dy <- liftIO $ readIORef dyRef
+    liftIO $ writeIORef dxyRef (Just (dx, dy))
+  return (xy, dxyRef)
+
+zipWithIncVar :: IncVar a -> IncVar b -> (a -> b -> c) -> (Delta a -> Delta b -> Delta c) -> IncM (IncVar c)
+zipWithIncVar x y f df = do
+  xy <- zipIncVar x y
+  fmapIncVar xy (uncurry f) (uncurry df)
+
+incZip2 :: IncVar a -> IncVar b -> IncM (IncVar (a, b))
+incZip2 x y = zipWithIncVar x y (,) (,)
+
+incZip3 :: IncVar a -> IncVar b -> IncVar c -> IncM (IncVar (a, b, c))
+incZip3 x y z = do
+  xy <- zipWithIncVar x y (,) (,)
+  zipWithIncVar xy z (\(a,b) c -> (a, b, c)) (\(a,b) c -> (a, b, c))
+
+instance (IncState a, IncState b, IncState c) => IncState (a, b, c) where
+  type Delta (a, b, c) = (Delta a, Delta b, Delta c)
+  applyDiff (x, y, z) (dx, dy, dz) = (applyDiff x dx, applyDiff y dy, applyDiff z dz)
+
+instance (IncState a, IncState b) => IncState (a, b) where
+  type Delta (a, b) = (Delta a, Delta b)
+  applyDiff (x, y) (dx, dy) = (applyDiff x dx, applyDiff y dy)
+
+
+addIncAction :: IO () -> IncM ()
+addIncAction action = modify \curAction -> curAction >> action
+
+-- === AllOrNothing class ===
+
+class (forall a. IncState (f a)) => AllOrNothing f where
+  fmapAllOrNothing :: IncVar (f a) -> (a -> b) -> IncM (IncVar (f b))
+
+instance AllOrNothing Unchanging where
+  fmapAllOrNothing v f = fmapIncVar v (\(Unchanging x) -> Unchanging (f x)) (const ())
+
+instance AllOrNothing Overwritable where
+  fmapAllOrNothing v f = fmapIncVar v (\(Overwritable x) -> Overwritable (f x)) (fmap f)
 
 -- === Delta type family ===