how can I decently add an "undo" functionality to State monads?

Question

Say that I have a State monad, and I want to do some manipulations on the state and might want to undo the change in future. How in general can I do this decently?

To give a concrete example, let's assume the state is just an Int, and the manipulation is just to increase the number by one.

type TestM a = StateT a IO ()

inc :: TestM Int
inc = modify (+ 1)

however, if I want to keep track of all the history of states in case I want to undo to some previous state, the best I can think of is to wrap the states in a stack: every modification to the state will be pushed to the stack so that I can undo changes through droping the top element on the stack.

-- just for showing what's going on
traceState :: (MonadIO m, MonadState s m, Show s) => m a -> m a
traceState m = get >>= liftIO . print >> m

recordDo :: TestM a -> TestM [a]
recordDo m = do
    x <- gets head
    y <- liftIO $ execStateT m x
    modify (y:)

inc' :: TestM [Int]
inc' = recordDo inc

undo' :: TestM [Int]
undo' = modify tail

-- inc 5 times, undo, and redo inc
manip' :: TestM [Int]
manip' = mapM_ traceState (replicate 5 inc' ++ [undo',inc'])

main :: IO ()
main = do
    v1 <- execStateT (replicateM_ 5 (traceState inc)) 2
    v2 <- execStateT (replicateM_ 5 (traceState inc')) [2]
    v3 <- execStateT manip' [2]
    print (v1,v2,v3)

As expected, here is the output:

2
3
4
5
6
[2]
[3,2]
[4,3,2]
[5,4,3,2]
[6,5,4,3,2]
[2]
[3,2]
[4,3,2]
[5,4,3,2]
[6,5,4,3,2]
[7,6,5,4,3,2]
[6,5,4,3,2]
(7,[7,6,5,4,3,2],[7,6,5,4,3,2])

The drawback of my approach:

• tail and head are unsafe
• One have to use something like recordDo explicitly, but I guess this is unavoidable because otherwise there will be some inconsistency issue. For example increasing the number by two can be done by either inc' >> inc' or recordDo (inc >> inc) and these two approach have different effects on the stack.

So I'm looking for either some ways to make it more decent or something that does the job of "reversible state" better.

Answer 1

Depending on your use-case, it might be worth considering something that I'd call "delimited undo":

{-# LANGUAGE FunctionalDependencies, FlexibleContexts #-}
import Control.Applicative
import Control.Monad
import Control.Monad.State
import Control.Monad.Trans.Maybe

undo :: (MonadState s m, MonadPlus m) => m a -> m a -> m a
undo dflt k = do
    s <- get
    k `mplus` (put s >> dflt)

undoMaybe :: (MonadState s m) => m a -> MaybeT m a -> m a
undoMaybe dflt k = do
    s <- get
    r <- runMaybeT k
    maybe (put s >> dflt) return r

undoMaybe_ :: (MonadState s m) => MaybeT m () -> m ()
undoMaybe_ = undoMaybe (return ())

Executing undo x k means "execute k, and if it fails, undo the state and execute x instead". Function undoMaybe works similarly, but allows the failure only the nested block. Your example then could be expressed as:

type TestM a = StateT a IO ()

inc :: (MonadState Int m) => m ()
inc = modify (+ 1)

-- just for showing what's going on
traceState :: (MonadIO m, MonadState s m, Show s) => m a -> m a
traceState m = get >>= liftIO . print >> m

inc' :: (MonadIO m, MonadState Int m) => m ()
inc' = traceState inc

-- inc 5 times, undo, and redo inc
manip' :: TestM Int
manip' = replicateM 4 inc' >> undoMaybe_ (inc' >> traceState mzero) >> inc'

main :: IO ()
main = do
    v1 <- execStateT (replicateM_ 5 (traceState inc)) 2
    putStrLn ""
    v3 <- execStateT manip' 2
    print (v1,v3)

The main advantage is that you can never underflow the stack. The disadvantage is that you can't access the stack and the undo is always delimited.

One could also create an Undo monad transformer that where the above undo becomes mplus. Whenever a failed computation is restored with mplus, the state is restored as well.

newtype Undo m a = Undo (m a)
    deriving (Functor, Applicative, Monad)

instance MonadTrans Undo where
    lift = Undo

instance (MonadState s m) => MonadState s (Undo m) where
    get = lift get
    put = lift . put
    state = lift . state

instance (MonadPlus m, MonadState s m) => MonadPlus (Undo m) where
    mzero = lift mzero
    x `mplus` y = do
        s <- get
        x `mplus` (put s >> y)