checkpoint.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

{-# OPTIONS_GHC -Wall #-}

module Checkpoint
  ( FreshT (..)
  , runFreshT
  , MonadFresh (..)

  , ValueThunk (..)
  , Checkpoints (..)
  , emptyCheckpoints
  , insertCheckpoint
  , lookupCheckpoint

  , CheckpointT (..)
  , runCheckpointT
  , MonadCheckpoint (..)

  , demo
  )
where

import Control.Monad.State (MonadState (..), StateT (..), evalStateT, runStateT)
import Control.Monad.Trans.Class (MonadTrans (..))
import Data.Aeson (FromJSON (..), ToJSON (..), Value)
import Data.Aeson.Types (parseMaybe)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (..))
import Data.Tuple (swap)
import Type.Reflection (Typeable, pattern TypeRep, (:~~:) (..), eqTypeRep, typeOf)

newtype FreshT m a = FreshT{ unFreshT :: StateT Word m a }
  deriving newtype (Functor, Applicative, Monad, MonadTrans)

runFreshT :: Monad m => FreshT m a -> m a
runFreshT (FreshT m) = evalStateT m 0

class MonadFresh m where
  fresh :: m Word

instance Monad m => MonadFresh (FreshT m) where
  fresh = do
    n <- FreshT get
    FreshT $ put (n + 1)
    pure n

data ValueThunk where
  Haskell :: (ToJSON a, Typeable a) => a -> ValueThunk
  JSON :: Value -> ValueThunk

instance ToJSON ValueThunk where
  toJSON = \case
    Haskell a -> toJSON a
    JSON v -> v

instance FromJSON ValueThunk where
  parseJSON v = pure (JSON v)

newtype Checkpoints = Checkpoints (Map Word ValueThunk)
  deriving newtype (ToJSON, FromJSON)

emptyCheckpoints :: Checkpoints
emptyCheckpoints = Checkpoints Map.empty

insertCheckpoint :: (ToJSON a, Typeable a) => Word -> a -> Checkpoints -> Checkpoints
insertCheckpoint n a (Checkpoints cs) = Checkpoints $ Map.insert n (Haskell a) cs

lookupCheckpoint :: (FromJSON a, Typeable a) => Word -> Proxy a -> Checkpoints -> Maybe a
lookupCheckpoint n (_ :: Proxy a) (Checkpoints cs) = do
  c <- Map.lookup n cs
  case c of
    Haskell a -> do
      HRefl <- eqTypeRep (TypeRep @a) (typeOf a)
      pure a
    JSON v -> parseMaybe parseJSON v

newtype CheckpointT m a = CheckpointT{ unCheckpointT :: StateT Checkpoints (FreshT m) a }
  deriving newtype (Functor, Applicative, Monad)

instance MonadTrans (CheckpointT) where
  lift = CheckpointT . lift . lift

runCheckpointT :: Monad m => Checkpoints -> CheckpointT m a -> m (Checkpoints, a)
runCheckpointT s (CheckpointT m) = fmap swap . runFreshT . flip runStateT s $ m

class MonadCheckpoint m where
  checkpoint :: (ToJSON a, FromJSON a, Typeable a) => m a -> m a

instance Monad m => MonadCheckpoint (CheckpointT m) where
  checkpoint :: forall a. (ToJSON a, FromJSON a, Typeable a) => CheckpointT m a -> CheckpointT m a
  checkpoint m = do
    cs <- CheckpointT get
    n <- CheckpointT . lift $ fresh
    case lookupCheckpoint n (Proxy @a) cs of
      Nothing -> do
        a <- m
        CheckpointT $ put (insertCheckpoint n a cs)
        pure a
      Just a -> do
        pure a

demo :: IO ()
demo = do
  let program :: CheckpointT IO Int
      program = do
        x <- checkpoint $ lift $ putStrLn "|> line 1" *> pure 1
        y <-              lift $ putStrLn "|> line 2" *> pure 2
        z <- checkpoint $ lift $ putStrLn "|> line 3" *> pure 3
        pure $ x + y + z

  putStrLn "--- 1st run (from scratch) ---"
  (cs, x) <- runCheckpointT emptyCheckpoints program
  putStrLn $ "Result: " <> show x

  putStrLn "--- 2nd run (with in-memory checkpoints) ---"
  (_, y) <- runCheckpointT cs program
  putStrLn $ "Result: " <> show y

  putStrLn "--- 3rd run (with JSON checkpoints) ---"
  let cs' = fromMaybe (error "unreachable") (parseMaybe parseJSON (toJSON cs))
  (_, z) <- runCheckpointT cs' program
  putStrLn $ "Result: " <> show z

-- ghci> demo
-- --- 1st run (from scratch) ---
-- |> line 1
-- |> line 2
-- |> line 3                                                      Result: 6
-- --- 2nd run (with in-memory checkpoints) ---
-- |> line 2
-- Result: 6
-- --- 3rd run (with JSON checkpoints) ---
-- |> line 2                                                      Result: 6