path: root/src/lib/IncState.hs

-- Copyright 2023 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE UndecidableInstances #-}

module IncState (
  IncState (..), MapEltUpdate (..), MapUpdate (..),
  Overwrite (..), TailUpdate (..), Unchanging (..), Overwritable (..),
  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 ===

class Monoid (Delta s) => IncState s where
  type Delta s :: *
  applyDiff :: s -> Delta s -> s

-- === Diff utils ===

data MapEltUpdate s =
   Create s
 | Replace s  -- TODO: should we merge Create/Replace?
 | Update (Delta s)
 | Delete
   deriving (Generic)

newtype MapUpdate k s = MapUpdate { mapUpdates :: M.Map k (MapEltUpdate s) }

mapUpdateMapWithKey
  :: (IncState s, IncState s')
  => MapUpdate k s -> (k -> s -> s') -> (k -> Delta s -> Delta s') -> MapUpdate k s'
mapUpdateMapWithKey (MapUpdate m) fs fd =
  MapUpdate $ flip M.mapWithKey m \k v -> case v of
    Create  s -> Create $ fs k s
    Replace s -> Replace $ fs k s
    Update d  -> Update $ fd k d
    Delete    -> Delete

instance (IncState s, Ord k) => Monoid (MapUpdate k s) where
  mempty = MapUpdate mempty

instance (IncState s, Ord k) => Semigroup (MapUpdate k s) where
  MapUpdate m1 <> MapUpdate m2 = MapUpdate $
    M.mapMaybe id (M.intersectionWith combineElts m1 m2)
      <> M.difference m1 m2
      <> M.difference m2 m1
    where combineElts e1 e2 = case e1 of
            Create s -> case e2 of
              Create _ -> error "shouldn't be creating a node that already exists"
              Replace s' -> Just $ Create s'
              Update d -> Just $ Create (applyDiff s d)
              Delete   -> Nothing
            Replace s -> case e2 of
              Create _ -> error "shouldn't be creating a node that already exists"
              Replace s' -> Just $ Replace s'
              Update d -> Just $ Replace (applyDiff s d)
              Delete   -> Nothing
            Update d -> case e2 of
              Create _ -> error "shouldn't be creating a node that already exists"
              Replace s -> Just $ Replace s
              Update d' -> Just $ Update (d <> d')
              Delete   -> Just $ Delete
            Delete -> case e2 of
              Create s -> Just $ Replace s
              Replace _ -> error "shouldn't be replacing a node that doesn't exist"
              Update _  -> error "shouldn't be updating a node that doesn't exist"
              Delete    -> error "shouldn't be deleting a node that doesn't exist"

instance (IncState s, Ord k) => IncState (M.Map k s) where
  type Delta (M.Map k s) = MapUpdate k s
  applyDiff m (MapUpdate updates) =
          M.mapMaybe id (M.intersectionWith applyEltUpdate m updates)
       <> M.difference m updates
       <> M.mapMaybe applyEltCreation (M.difference updates m)
    where applyEltUpdate s = \case
            Create _ -> error "key already exists"
            Replace s' -> Just s'
            Update d -> Just $ applyDiff s d
            Delete   -> Nothing
          applyEltCreation = \case
            Create s -> Just s
            Replace _ -> error "key doesn't exist yet"
            Update _  -> error "key doesn't exist yet"
            Delete    -> error "key doesn't exist yet"

data TailUpdate a = TailUpdate
  { numDropped :: Int
  , newTail    :: [a] }
  deriving (Show, Generic)

instance Semigroup (TailUpdate a) where
  TailUpdate n1 xs1 <> TailUpdate n2 xs2 =
    let xs1Rem = length xs1 - n2 in
    if xs1Rem >= 0
      then TailUpdate n1 (take xs1Rem xs1 <> xs2) -- n2 clobbered by xs1
      else TailUpdate (n1 - xs1Rem) xs2           -- xs1 clobbered by n2

instance Monoid (TailUpdate a) where
  mempty = TailUpdate 0 []

instance IncState [a] where
  type Delta [a] = TailUpdate a
  applyDiff xs (TailUpdate numDrop ys) = take (length xs - numDrop) xs <> ys

-- Trivial diff that works for any type - just replace the old value with a completely new one.
data Overwrite a = NoChange | OverwriteWith a
  deriving (Show, Eq, Generic, Functor, Foldable, Traversable)
newtype Overwritable a = Overwritable { fromOverwritable :: a } deriving (Show, Eq, Ord)

instance Semigroup (Overwrite a) where
  l <> r = case r of
    OverwriteWith r' -> OverwriteWith r'
    NoChange         -> l

instance Monoid (Overwrite a) where
  mempty = NoChange

instance IncState (Overwritable a) where
  type Delta (Overwritable a) = Overwrite a
  applyDiff s = \case
    NoChange         -> s
    OverwriteWith s' -> Overwritable s'

-- Case when the diff and the state are the same
newtype MonoidState a = MonoidState a

instance Monoid a => IncState (MonoidState a) where
  type Delta (MonoidState a) = a
  applyDiff (MonoidState d) d' = MonoidState $ d <> d'

-- Trivial diff that works for any type - just replace the old value with a completely new one.
newtype Unchanging a = Unchanging { fromUnchanging :: a } deriving (Show, Eq, Ord)

instance IncState (Unchanging a) where
  type Delta (Unchanging a) = ()
  applyDiff s () = s

instance            ToJSON a  => ToJSON (Overwrite a)
instance (ToJSON k, ToJSON s, ToJSON (Delta s)) => ToJSON (MapUpdate k s) where
  toJSON m = toJSON $ M.toList $ mapUpdates m
instance ToJSON a => ToJSON (TailUpdate a)
instance (ToJSON s, ToJSON (Delta s)) => ToJSON (MapEltUpdate s)