Feeding stateful and stateful-effectful computations into each other

Question

I am writing code that runs a simulation in an interactive way. The building blocks are:

runStep :: I -> State S O
runWorld :: World m -> O -> StateT W m I

Note that the world has some state W but is also allowed to have effects of type m. My question is about the function that feeds runStep and runWorld into each other. My current implementation is quite heavy on the Control.Lens noise:

sim :: (Monad m) => World m -> StateT (I, S, O) m Result
sim w = do
  inp <- use _1
  out <- zoom _2 $ state . runState $ runStep inp
  inp' <- zoom _3 $ runWorld w out
  _1 .= inp'
  return $ resultComputedPurelyFrom out

What I'd like to improve on this code is two things:

• Reduce the lens piping. I could, of course, get rid of _1, _2 and _3 by using a bespoke record type, but then that introduces its own noise -- sim is the only function that uses this particular combination of I, S and O.

• The write-back to inp' could easily be accidentally omitted, resulting in not-obviously-wrong-looking, well-typed code that wouldn't work correctly.

ETA: An alternative formulation is to forego the lens stuff completely, at the cost of making even the updating of the S component something that can be accidentally omitted:

sim :: (Monad m) => World m -> StateT (I, S) (StateT O m) Result
sim w = do
  (inp, s) <- get 
  let (out, s') = runState (runStep inp) s
  inp' <- lift $ runWorld w out
  put (inp', s')
  return $ resultComputedPurelyFrom out

Answer 1

What I'd like to improve on this code is two things:

• Reduce the lens piping. I could, of course, get rid of _1, _2 and _3 by using a bespoke record type, but then that introduces its own noise -- sim is the only function that uses this particular combination of I, S and O.

Is this really a problem? I don't understand what could be improved about that part.

• The write-back to inp' could easily be accidentally omitted, resulting in not-obviously-wrong-looking, well-typed code that wouldn't work correctly.

You can move the state-passing of I into a separate function for which that's the only job:

withInput :: ALens' s i -> (i -> StateT s i) -> StateT s ()
withInput l go = do
  i <- use (cloneLens l)
  i' <- go i
  cloneLens l := i'

...
sim = do
  withInput _1 $ \inp ->
    inp' <- ...
    pure inp'
  ...

The result type () of withInput makes it difficult to forget to update the state: the function would fail to typecheck if you forget the last clause.