Strictly compute and memoize in discriminated unions, in F#

Question

I have a discriminated union, such as

type Dish = 
| Eggs
| Spam of Dish

This is basically a linked list, without any content, e.g. Spam(Spam(Spam(Eggs))). I want to strictly perform a computation on this structure, such as counting the length, and memorize the result. In a normal type, I'd use class-local let bindings, but those aren't available in discriminated unions.

One way to do this would be,

type Count = int
type Dish = 
| Eggs
| Spam of Dish * Count

But this is really messy, when the data I need is easily computable, but I still hope there is a better way (without using external mutable constructs).

Answer 1

After reviewing the answers, I've decided to go with a model that seems the least obtrusive to me. I've used a modified object to demonstrate how it would work in a slightly more complex scenario.

type StackDef<'a>(v : 'a, s : Stack<'a>) = 
    member val Length = s.Length + 1
    member val Inner = v, s 
and Stack<'a> = 
    | Empty
    | Stack of StackDef<'a>
    member this.Length = 
        match this with
        | Empty -> 0
        | Stack(def) -> def.Length

let Stack (v, s) = Stack(StackDef(v, s)) 
let (|Stack|Empty|) = function | Empty -> Empty | Stack(sd) -> Stack(sd.Inner)
//...
let example = Stack(1, Stack(2, Stack(3, Empty))).Length

1. It doesn't contain any external mutable state.
2. The discriminated union Dish (or in the example, Stack) continues to exist.
3. The field length doesn't appear in the union definition at all, nor is it provided by any constructor, just as it should be.
4. The memoized data is associated with the instance, as it should be.

However, having thought about it, by using a static weaver such as Afterthought it might be possible to replace any method such as:

Stack<'a> = 
    | Empty
    | Stack of 'a * Stack<'a>

    [<Lazy>] //custom attribute that would work with a static weaver
    member this.Length = 
        match this with
        | Empty -> 0
        | Stack(_, s) -> s.Length + 1

With a private readonly Lazy<int> __length initialized in the constructor with a delegate that executes the above code, and change the actual content of the method to simply invoking __length.Value. While F# doesn't allow union types to contain fields, possibly for very valid reasons, I highly doubt the IL would have such restrictions.
In fact, it would be possible to do a lot of things using some IL manipulation. Maybe it's something to think about.

Answer 2

One option is making the union cases private to hide the cached length.

//the 'guts' of Dish -- entirely hidden
type private DishImpl = 
  | Eggs
  | Spam of DishImpl

// Dish wrapper type -- implementation hidden
type Dish = 
  private 
  | Dish of DishImpl * int
  with
    // O(1), just get the 'length' field
    member x.Length = let (Dish(_, len)) = x in len
    static member Eggs() = Dish(Eggs, 1)
    static member Spam(Dish(dish, len)) = Dish(Spam dish, len + 1)

let eggs = Dish.Eggs()
let spam = Dish.Spam(eggs)
printfn "%d" eggs.Length //outputs: 1
printfn "%d" spam.Length //outputs: 2

To do it up right, create an accompanying module with let-bound functions and active patterns for destructuring.

Answer 3

If you tolerate a bit internal mutable state, here is a memoize function which creates a dictionary per function:

let memoize f =
    let dict = Dictionary()
    fun n ->
        match dict.TryGetValue(n) with
        | (true, v) -> v
        | _ ->
            let res = f(n)
            dict.Add(n, res)
            res
// This function results in a warning though
let rec length = memoize (function Eggs -> 0 | Spam d -> 1 + length d)

The approach isn't that bad since the mutable dictionary is hidden.

A purely functional approach could be using Map to hold values and a kind of State computation expression to hide Map values passing around. Please refer to this snippet to see how a memoize computation expression looks like.

Answer 4

There is also Memo Functions, Polytypically! by Ralph Hinze (2000). Adapting to F#:

type Dish =
    | Eggs
    | Spam of Dish

type DishTable<'T> =
    {
        Eggs : Lazy<'T>
        Spam : Lazy<DishTable<'T>>
    }

let rec tabulate (f: Dish -> 'T) : DishTable<'T> =
    {
        Eggs = lazy f Eggs
        Spam = lazy tabulate (f << Spam)
    }

let rec lookup (table: DishTable<'T>) (dish: Dish) : 'T =
    match dish with
    | Eggs -> table.Eggs.Value
    | Spam x -> lookup table.Spam.Value x

let memo (f: Dish -> 'T) : (Dish -> 'T) =
    lookup (tabulate f)

let rec len x =
    match x with
    | Eggs -> 0
    | Spam x -> 1 + len x

let l2 = memo len

Answer 5

This is what I came up with. It's not true memoization because it counts eagerly when you call mem, but might work for your needs.

type Dish = 
| Eggs
| Spam of Dish 
| Memo of Dish * int
with 
    member d.length = 
        match d with 
        | Eggs        -> 1
        | Spam d      -> 1 + d.length
        | Memo (d, i) -> i
    member d.mem = 
        match d with 
        | Eggs        -> Memo(d, d.length)
        | Spam d2     -> Memo(d, d.length)
        | Memo(d2, i) -> d // no need to memo it again

let x = Spam (Spam(Spam Eggs))
let m = x.mem

x.length // val it : int = 4
m.length // val it : int = 4

Answer 6

Note that in your case, literally the only interesting property of a value of your type is its length, so you might as well just use integers as your representation instead:

let Eggs = 0
let Spam n = 1 + n

let (|Eggs|Spam|) = function
| 0 -> Eggs
| n -> Spam(n-1)

let length = id

// example usage
let dish = Spam(Spam(Eggs))

let l = length dish

let kind =
    match dish with
    | Eggs -> "Eggs"
    | Spam(Eggs) -> "One Spam"
    | Spam(Spam _) -> "At least two Spams"

If your real question is how to do this for a more interesting type, then one approach would be to create mutually recursive types, one of which is annotated:

type 'a AnnotatedDish = { dish : 'a Dish; value : 'a }
and 'a Dish =
| Eggs
| Spam of 'a AnnotatedDish

// "smart" constructors, given that you want to annotate with length
let eggs = { dish = Eggs; value = 0 }
let spam d = { dish = Spam d; value = 1 + d.value }

let length { value = l } : int = l

// active patterns
let (|Eggs|Spam|) = function
| { dish = Eggs } -> Eggs
| { dish = Spam d } -> Spam d


// example usage
let dish = spam(spam(eggs))

let l = length dish

let kind =
    match dish with
    | Eggs -> "Eggs"
    | Spam(Eggs) -> "One Spam"
    | Spam(Spam _) -> "At least two Spams"