Writer and State Monad in Rust, based on higher

Motivation

As can for instance be seen in my post about Free Monads, working with Monads in Rust is not as straightforward as it is in languages with higher kinded types. However, the relative success with Free Monads made me start a spare time game project that uses a Free Monad based eDSL to describe the story flow. I am planning to use the do-notation offered by the run!{} macro from the "higher" crate to build the story, but I've found that the run!{} macro has a rather annoying problem, namely that bindings created in it only really work for types that are Copy. While a pull request I made to add explicit clone support has been merged, that is only a bandaid with terrible ergonomics.

However, even with Copy types, it would be quite annoying to have to pass the game state along manually all the time. This is where the State Monad Transformer comes in. Within do-notation, it almost feels like mutable global state, even though nothing is global, and all computations are pure.

In addition, I want to record all player choices since the last save-point. This is meant as a way to allow saving even in situations that normally would not. This sounds like a use case for a Writer Monad Transformer, as building a list of things from within do-notation is one of its primary use cases.

The Writer Monad Transformer

I started out with the Writer Monad Transformer first, because that was the first one I had thought about. My spare time project was at a point where the next logical step was to add savegame support, and, as said, I thought the Writer Monad Transformer could help me there.

An attempt was made: How to (not) implement a Writer Monad Transformer in Rust

I usually prefer simple over complex code, as this makes it easier to discuss it with others, and – just as important – get back to the code at some point in the future. That's why I at first attempted to directly port Haskell's Lazy Writer Monad Transformer. (For the remainder of this blog post, if I talk about the "Lazy Writer Monad Transformer", I mean the implementation that is lazy in Haskell - the Rust code shown below is very much eager.) Unlike the CPS implementation of it, the lazy version is straightforward. It takes the wrapped Monad, and replaces its contents with a pair of the original content and an accumulator. Nothing special about it, right?

Turns out, it's not that easy. Implementing Functor for it is straightforward enough:
struct WriterT<M>{
  run_writer_t : M
}

impl<'a, A, M, W> Functor<'a, A> for WriterT<M> where M : Functor<'a, (A,W)>{
  type Target<T> = WriterT<M::Target<(T, W)>>;
  fn map<F>(self, f : F) -> Self::Target<B> where F : Fn(A)->B{
    let f = Rc::new(f); //lifetimes
    WriterT{ run_writer_t :
      self.run_writer_t.fmap(|(a,w)| (f(a), w))
    }
  }
}

However, to be useful in do-notation, the type also needs to offer Bind and Pure. Here, however, a problem appears that (again) would require Non-Lifetime Binders. While there are many ways to express the required operation, they all need to somehow call a method on a GAT. For instance, it could look something like this:
impl<'a, A, M ,W> Bind<'a, A> for WriterT<M> where M : Bind<'a, (A,W)>, M : Clone, ???? {
  type Target<T> = WriterT<M::Target<(T, W)>>;
  fn bind<F>(self, f : F) -> Self::Target<B> where F : Fn(A)->Self::Target<B> {
    let f = Rc::new(f); //lifetimes...
    self.run_writer_t.bind(move |(a,w)| f(a).fmap(move |(b,ww)| (b, w.clone().mappend(ww))))
  }
}
I'm writing this from memory, so this might not compile. In any case, it won't compile because of the missing trait bound, denoted by ????. What it would need to express would be that this trait implementation should only exist if f(a) is a Functor. Adding that trait bound to the bind method itself would be trivial. Adding it to the whole trait implementation would require Non-Lifetime binders. In this formulation it would even require a more advanced version of Non-Lifetime Binders, that would not only allow to express "for all types B", but rather "for all types B that are Sized" due to the implicit Sized bound on GATs.

I was about to give up on a generic implementation at this point, but then I remembered that I had skimmed through an alternative implementation of the Writer Monad Transformer, that had a Bind implementation which has easier to express trait bounds, the CPS Writer Monad Transformer. Originally I had ignored this version because it is a bit more complex than the lazy implementation, but given that the lazy implementation seems to be a dead end (in Stable Rust), I just had to give it a try.

Continuation Passing Style Writer Monad Transformer

Looking over Haskell's source code for the CPS Writer Monad Transformer made me quite optimistic. Apply would likely not be implementable in a generic way, again because it would require Non-Lifetime Binders, but unlike Bind, Apply is not required to use do-notation. So, while the resulting type could not implement the Apply, Applicative, and, ultimately, the Monad type trait, that's not a problem for my use case. Also, technically, it would still qualify as a Monad. The requirement that Monad is a subtrait of Apply in higher is for historical, not technical reasons.

The drawback of the CPS version of WriterT is that it's lazy (this time also in Rust, not only Haskell). In addition, if the code should really be generic, it cannot make any assumptions about lifetimes, meaning that there will be boxed trait objects, shared references, and a lot of calls to clone(). Still, I was excited. In my use case the Writer was not going to end up in any hot path of the program. A few heap allocations and dynamic dispatch are quite an acceptable price at that point.

The main building block for a Continuation-Passing Writer Monad Transformer is a data type that wraps a function, which maps an accumulator to a Monad that wraps a pair of said accumulator and another value. Haskell is rather explicit here:
newtype WriterT w m a = WriterT { unWriterT :: w -> m (a, w) }
In Rust we cannot express the term m (a, w) due to a lack of higher-kinded types. The best we can do is just M:
struct WriterT<'a, W, M>{ un_writer_t : Box<dyn Fn(W) -> M + 'a> }
I don't think that making this a boxed trait object is really necessary, it could probably just be a generic parameter.

With this type, we can implement Functor, Bind and Pure rather easily. They only need to argue about the wrapped Monad M. Bind works, because instead of accessing the accumulator from the function call result, it can, in this formulation, simply pass the old accumulator as a parameter to it:
impl<'a, A, W : 'a, M : 'a> Functor<'a,A> for WriterT<'a, W,M> where M : Functor<'a, (A, W)>{
  type Target<T> = WriterT<'a, W, M::Target<(T,W)>>;
  fn fmap<B, F>(self, f: F) -> Self::Target<B>
  where
    F: Fn(A) -> B + 'a {
    let f = Rc::new(f);
    WriterT {
      un_writer_t: Box::new(move |w|{ let f = f.clone(); (self.un_writer_t)(w).fmap(move |(a, ww)| (f(a),ww))}),
    }
  }
}

impl<'a, A, W : 'a, M : 'a> Bind<'a, A> for WriterT<'a, W, M> where M : Bind<'a, (A, W)>{
  type Target<T> = WriterT<'a, W, M::Target<(T,W)>>;
  fn bind<B, F>(self, f: F) -> Self::Target<B>
  where
    F: Fn(A) -> Self::Target<B> + 'a {
    let f = Rc::new(f);
    WriterT{
      un_writer_t: Box::new(move |w|{
        let f = f.clone();
        run!{
          old <= (self.un_writer_t)(w);
          (f(old.0).un_writer_t)(old.1)
        }
      })
    }
  }
}

impl<'a, A : 'a, W: 'a, M: 'a> Pure<A> for WriterT<'a, W, M> where M : Pure<(A,W)>, A : Clone{
  fn pure(value: A) -> Self {
    WriterT { un_writer_t: Box::new(move |w| M::pure((value.clone(),w))) }
  }
}

Luckily, the functions to create a new WriterT, to unwrap the wrapped Monad, to add data to the accumulator, and to lift a command into the wrapped monad,can all be implemented as well:
impl<'a, W: 'a, M: 'a> WriterT<'a, W, M>{
  fn writer_t<A>(f : M) -> Self where M : Functor<'a,(A,W), Target<(A, W)> = M> + Clone, W : Monoid + Clone{
    WriterT { un_writer_t: Box::new(move |w| {
      f.clone().fmap(move |(a, ww)| { (a, w.clone().mappend(ww)) })
    }) }
  }
  fn run_writer_t(&self) -> M where W : Default{
    (self.un_writer_t)(W::default())
  }
  fn lift<N: 'a, A>(m : N) -> Self where N : Functor<'a, A, Target<(A,W)> = M> + Clone, W : Clone{
    WriterT { un_writer_t: Box::new(move |w|{
      m.clone().fmap(move |a| (a,w.clone()))
    }) }
  }
  fn writer<A : 'a>((a,ww) : (A, W)) -> Self where M : Pure<(A,W)>, W : Monoid + Clone, A : Clone {
    WriterT { un_writer_t: Box::new(move |w| M::pure((a.clone(), w.mappend(ww.clone())))) }
  }
  fn tell(w : W) -> Self where M : Pure<((),W)>, W : Monoid + Clone{
    Self::writer(((),w))
  }
}

The rest of the usual methods can be implemented as well, I won't repeat them here for brevity. The full source code with a small usage example can be found on this Playground. The code quality isn't too great, as this is more of a proof-of-concept than anything else. The Playground also contains the Writer Monad, which is just the Writer Monad Transformer wrapping the Identity Monad.

But wait, can I actually use this?

Now that the Writer Monad Transformer had been implemented, I realized something. It stores its accumulator in the Monad that it wraps. For a Free Monad this means, that the accumulator is stored within the Pure nodes, which form the leaves in the tree-like structure that the Free Monad forms. This means, that, if you get handed a Writer Monad Transformer that wraps a Free Monad, you need to traverse the whole structure, until the very end, in order to read the accumulator.

You can already guess where this is going. In the context of a Free Monad based eDSL for a game story, this means that the player needs to finish the game before the interpreter for the eDSL can actually access the accumulator, what makes it a bit unsuitable for a save-game system...

While this setback made me rethink which part of my spare time project should actually be responsible for recording user input, and made me in the end decide against making this a part of the eDSL and rather just let the interpreter do it, I am pretty sure that the idea would still be salvagable, if one really wanted to. The problem is that a Writer Monad Transformer that wraps a Free Monad does not yield the accumulator to an interpreter for intermediate steps. However, I think that a Free Monad Transformer that wraps a Writer Monad would.

The State Monad Transformer

The cool thing about the Continuation-Passing-Style Writer Monad Transformer is, that it is nearly identical to the State Monad Transformer. The difference is that the State Monad Transformer does not accumulate data into a Monoid, but instead gives read/write access to a single element.

The similarities go so far, that Functor, Pure and Bind can be implemented identically. I won't repeat them here, just take the ones from the Writer Monad Transformer above, and replace un_writer_t by run_state_t.

The lift associated function is different though, as for StateT it just replaces the stored data, instead of appending to it. Also, StateT has functions to put and get the state.

struct StateT<'a, S, M>{ run_state_t : Box<dyn Fn(S)->M + 'a> }
impl<'a, S: 'a, M: 'a> StateT<'a, S, M>{
  fn lift<N: 'a, A>(m : N) -> Self where N : Functor<'a, A, Target<(A,S)> = M> + Clone, S : Clone{
    StateT { run_state_t: Box::new(move |s|{
      m.clone().fmap(move |a| (a,s.clone()))
    }) }
  }
  fn get() -> Self where M : Pure<(S,S)>, S : Clone{
    Self::state(|s| (s.clone(),s))
  }
  fn put(s : S) -> Self where M: Pure<((), S)>, S : Clone{
    Self::state(move |_| ((),s.clone()))
  }
}

Of course there are more methods that can be implemented for the State Monad Transformer, these are just the bare minimum. You can have a look at this Playground to see a few additional methods.

But what about Apply?

The problem with Apply is that it requires a means to access the to-be-applied function. If we look at higher's Apply trait, we can see that the type passed to the apply function is <Self as Apply<'a, A>>::Target<ApplyFn<'a, A, B>>. To get the contained ApplyFn, we would need to convince the compiler that this type is at least a Functor. However, we cannot argue about this type with the compiler when implementing Apply, because it depends on the type of ApplyFn passed to it by the caller. It's only available via a GAT, and does not exist in the where clause that we would need to write - similar to the problem encountered with Bind for the Lazy Writer Monad Transformer. Again, Non-Lifetime Binders could solve this problem. For now, we can implement Apply for a specialized State/Writer Monad Transformer though. For a known Monad Type, the compiler can actually check which traits <Self as Apply<'a, A>>::Target<ApplyFn<'a, A, B>> implements. For instance, for the type OptionalWriter<'a, W, A> = WriterT<'a, W, Option<A>>; the implementation is straightforward:

impl<'a,A : 'a, W : 'a> Apply<'a,A> for OptionalWriter<'a,W,(A,W)> where Self: Clone{
  type Target<T> = OptionalWriter<'a,W,(T,W)>
  where
    T: 'a,
    A: 'a;
  fn apply<B>(
    self,
    f: <Self as Apply<'a, A>>::Target<higher::apply::ApplyFn<'a, A, B>>,
  ) -> <Self as Apply<'a, A>>::Target<B>
  where
    B: 'a {
    OptionalWriter{ un_writer_t: Box::new(move |w| {
      let u = self.clone().un_writer_t;
      (f.un_writer_t)(w).bind(move |(f,ww)| (u)(ww).fmap(move |(x,www)| (f.apply(x),www)))
    })}
  }
}

The where clause does not even mention <Self as Apply<'a, A>>::Target<ApplyFn<'a, A, B>> at all. The reason why this works, and the generic version does not, is that Rust does not enforce any constraints on GATs. The result of Functor::fmap() need not be a Functor in general. For concrete Monad types the compiler can however check if this is the case, and if yes, use this information.

Can this pattern be used elsewhere?

Certainly. The Reader Monad Transformer comes to mind, obviously, since it can be implemented as a State Monad Transformer that simply doesn't allow to write and starts with an initial state. Other types that have similar structure can follow this pattern as well.

Given that I did not yet manage to write a generic Free Monad in Stable Rust, I of course also checked if there is an implementation that follows a similar pattern. However, the closest I found was the Church Encoded Free Monad, which requires the ability to compose generic functions, something that I don't think Rust supports. I'm still thinking about this problem, but haven't found a nice solution yet.

Will this be on Crates.io?

I'm not happy with the code quality that came out of this. So, no, I'm not planning to upload these types to crates.io. At least not any time soon. If you want to use them, just grab the two playground links above, and copy&paste their contents. For documentation, check the Haskell docs for WriterT and StateT.


Previous: Upgrading from Zen 1 to Zen 3 with a B350M Mortar mainboard Home