State Monad For The Rest Of Us - Part 12

Arialdo Martini — 1/07/2024 — F# Functional Programming

Index
Source code: github.com/arialdomartini/state-monad-for-the-rest-of-us.

Here’s what you got in Chapter 11:

let rec index =
    function
    | Leaf(v: string) ->
        getCount
        >>= (fun count ->
            let leaf = Leaf(v, count)
            putCount (count + 1)
            >>= (fun _ -> pure' leaf))
    | Node(l, r) ->
        index l >>= (fun ll ->
            index r
             >>= (fun rr ->
                 pure' (buildNode ll rr)))

Remember again what >>= does:

  • You feed it with a WithCount v.
  • It continues the execution with a lambda just taking a naked v.

>>= is that magic function that moves WithCount out of sight, while still accounting for it under the scenes.

In fact, the source of mess is the necessary presence of that lambda. If you had a function foo returning a WithCount String:

// () -> WithCount String
let foo () = WithCount "hello"

how beautiful would it be if you could replace:

foo() >>= fun (v -> ...)

with just:

let! v = foo()

with a special let! binding able to unwrap any WithCount value, and directly assigning the naked value "hello" to v?
If only you had that let! binding, the whole index function would drammatically simplify. Instead of:

let rec index =
    function
    | Leaf(v: string) ->
        getCount
        >>= (fun count ->
            let leaf = Leaf(v, count)
            putCount (count + 1)
            >>= (fun _ -> pure' leaf))
    | Node(l, r) ->
        index l >>= (fun ll ->
            index r
             >>= (fun rr ->
                 pure' (buildNode ll rr)))

you could operate the following transformations:

Original expression Special form
getCount >>= fun count -> let! count = getCount
putCount (count + 1) >>= fun _ -> let! _ = putCount (count + 1)
index l >>= fun ll -> let! ll = index l
index r >>= fun rr -> let! rr = index r

and the whole index function would look like:

let rec index =
    function
    | Leaf(v: string) ->
        let! count = getCount
        let leaf = Leaf(v, count)
        let! _ = putCount (count + 1)
        pure' leaf
    | Node(l, r) ->
        let! ll = index l
        let! rr = index r
        pure' (buildNode ll rr)

Oh, this would be really a game changer, wouldn’t it? It would allow you to think imperatively, while being purely functional.

Of course, if that let! binding existed, we would expect that under the hood the F# compiler would automatically convert any:

let! vv = v
...

into:

v >>= (fun vv -> ...)

Being a mechanical, deterministic transformation, that should not be an impossible dream.

Enter Computation Expressions

It turns out that F# does support that syntax. All you have to do is to teach it which >>= and pure' functions to use for those tranformations. You need to create a custom type and to implement a couple of methods:

type WithCountExpression() =
    member this.Return v = failwith "Not yet implemented" 
    member this.Bind v f = failwith "Not yet implemented" 

Return is an alias name for pure'. Bind, as you already know, is the name of >>=. So, after all, you should already know how to implement both:

type WithCountExpression() =
    member this.Return(v) = pure' v 
    member this.Bind(v, f) = v >>= f 

Now, just create an instance of it:

let withCount = WithCountExpression()

and, voilà!, you can use the new syntax:

let rec index =
    function
    | Leaf v ->
        withCount {
            let! count = getCount
            let leaf = Leaf (v, count)
            let! _ = putCount (count + 1)
            return leaf
        }
    | Node(l, r) ->
        withCount {
            let! ll = index l
            let! rr = index r
            return buildNode ll rr
        }

This is mostly what we initially desired, besides the extra withCount { } wrapping the whole epxression. The reason why this is needed is because you can have multiple monads, and for each of them let! and return must be transformed to a specific implementation of >>= and pure'. This is a problem Haskell does not have: it provides you with one single do expression, covering all the possible existing and future monads.

Computation Expression

What you just implemented is called Computation Expression. It is equivalent to (and possibly more powerful than) the Haskell do notation, which would allow writing index as:

index (Leaf v) =
    do {
        count <- getCount
        let leaf = Leaf (v, count)
        putCount (count + 1)
        return leaf
    }
index (Node l r) =
    do {
        ll <- index l
        rr <- index r
        return buildNode ll rr
    }

You see the similarity, I hope!

Computation Expressions are super powerful, and they are pervasive in F#: they can handle async code, error management, optional values, sequence generations, logging etc. After all, they are syntactic sugar around arbitrary monads.

More sugar, please!

A little trick: when you want to ignore a value, you can always replace let! _ = ... with do! .... So, you can write index as:

let rec index =
    function
    | Leaf v ->
        withCount {
            let! count = getCount
            let leaf = Leaf (v, count)
            do! putCount (count + 1)
            return leaf
        }
    | Node(l, r) ->
        withCount {
            let! ll = index l
            let! rr = index r
            return buildNode ll rr
        }

Naturally, it is again just syntactic sugar. Any:

do! foo()
...

is replaced with:

foo() >>= (fun _ -> ...)

It’s all fake imperative style

Let me stress that, although expressions like:

withCount {
  do! putCount 42
  do! putCount 0
  do! putCount 99
  return "Hello, world!"
}

look like a series of statements, under the hood each do! like is an expression *bound` to the next one: the whole is a unique expression, in a chain of lambdas:

putCount 42 >>= (fun _ ->
      putCount 0 >>= (fun _ ->
          putCount 99 >>= (fun _ ->
              pure' "Hello, world!")))

Whatever you write inside a withCount is a single, purely functional expression.

What about LINQ?

In the index I claimed that this chapter would help you see LINQ for what it is: a monadic engine. It you always saw LINQ as an embedded SQL language, consider this implementation of index, converted from F# to C#:

private static WithCount<Tree<(A, int)>> IndexLINQ<A>(Tree<A> tree) =>
    tree switch
    {
        Tree<A>.Leaf(var v) =>
            from count in GetCount
            from _ in PutCount(count + 1)
            select BuildLeaf(v, count),
            
        Tree<A>.Node(var l, var r) =>
            from ll in Index(l)
            from rr in Index(r)
            select BuildNode(ll, rr)
    };

It should not be hard to see these mappings:

F# syntax C# LINQ syntax
let! v = monadicFunction() from v in monadicFunction()
return v select v

Just like F#, in C# you can teach LINQ how to monadically handle WithCount values, instructing it which >>= / Bind methods to use. In C# the way to go is with Extension Methods:

internal static class WithCountExtensions
{
    internal static WithCount<T> Pure<T>(T value) => new(count => (value, count));

    internal static (T, int) Run<T>(WithCount<T> value, int count) => value.F(count);

    internal static WithCount<TResult> Bind<T, TResult>(WithCount<T> a, Func<T, WithCount<TResult>> f) =>
        new(count =>
        {
            var (va, ca) = Run(a, count);
            var result = f(va);
            return Run(result, ca);
        });

    internal static WithCount<TResult> SelectMany<T, TIntermediate, TResult>(
        this WithCount<T> withCount,
        Func<T, WithCount<TIntermediate>> intermediateSelector,
        Func<T, TIntermediate, TResult> resultSelector) =>
        new(count =>
        {
            var (value, intermediateCount) = withCount.F(count);
            var intermediateWithCount = intermediateSelector(value);
            var (intermediateValue, finalCount) = intermediateWithCount.F(intermediateCount);
            return (resultSelector(value, intermediateValue), finalCount);
        });
}

Yes, it’s way less readable than the equivalent F# version — now you see why functional programmers prefer F# over C#, don’t you?.
You can find the complete C# implementation in StateMonadTest.cs.

Anyway, here’s the take away: LINQ is much more than a tool for embedding SQL and manipulating lists. Just like the Haskell’s *do notation and F# Computation Expressions, LINQ is syntactic sugar for monads, any monad. LINQ is indeed a monad engine.

You can read more about it in Thinking Functionally: What is LINQ really?, from the language-ext’s wiki.

Much simpler than this

You made it to the end! Congrats! If you found getting to the State Monad arduous, it is because it really was.

Here’s something that might catch you off guard: working with a State Monad is really trivial.
I really think that this is a general trait of Functional Programming: using functional programming concepts is orders of magnitude simpler than re-implementing them from the scratch. What you have done, in these 12 strenuous chapters, was to distill a State Monad out of thin air, from the ground up, incrementally and taking many detours for reasoning about any little nuance as you stumbled upon along the journey.

In the next — much easier — chapters, I want to show you how easy and smooth working with a State Monad is, once it is implemented.

Stay tuned! I will publish them in the next days!

References

Comments

GitHub Discussions

Functional Programming For The Rest Os Us

Interested in FP? Be the first to be notified when new introductory articles on the topic are published.