Inlining imperative computation

Imperative.fs

module Imperative

open System

[<Struct>]
type BlockResult<'T> = NoValue | Return of 'T | Continue | Break with
    member this.Close() =
        match this with
        | NoValue -> ValueNone
        | Return x -> ValueSome x
        | Continue -> failwith "invalid continue"
        | Break -> failwith "invalid break"

type Code<'T> = Ref<BlockResult<'T>> -> unit

type BlockBuilder() =
    member inline this.Return x: Code<_> =
        fun res -> res.Value <- Return x

    member inline this.Delay ([<InlineIfLambda>] f): Code<_> =
        fun res -> (f()) res

    member inline this.Zero(): Code<_> = ignore

    member inline _.While([<InlineIfLambda>] cond : unit -> bool, [<InlineIfLambda>] body : Code<'T>) : Code<'T> =
        fun res ->
            let mutable looping = cond()
            while looping do
                body res
                match res.Value with
                | NoValue -> looping <- cond()
                | Return _ -> looping <- false
                | Break -> res.Value <- NoValue; looping <- false
                | Continue -> res.Value <- NoValue; looping <- cond()

    member inline _.Combine([<InlineIfLambda>] part1: Code<'T>, [<InlineIfLambda>] part2: Code<'T>): Code<'T> =
        fun res ->
            part1 res
            match res.Value with
            | NoValue -> part2 res
            | _ -> ()

    member inline _.YieldFrom(code) = code

    member inline _.TryFinally([<InlineIfLambda>] body: Code<'T>, [<InlineIfLambda>] compensation : unit -> unit) : Code<'T> =
        fun res ->
            try
                body res
            finally
                compensation()

    member inline b.Using(disp : #IDisposable, [<InlineIfLambda>] body: #IDisposable -> Code<'T>) : Code<'T> =
        // A using statement is just a try/finally with the finally block disposing if non-null.
        b.TryFinally(
            (fun res -> (body disp) res),
            (fun () ->
                match disp with
                | null -> ()
                | _ -> disp.Dispose()))

    member inline b.For(sequence: seq<'TElement>, [<InlineIfLambda>] body: 'TElement -> Code<'T>) : Code<'T> =
        b.Using (sequence.GetEnumerator(),
            (fun e -> b.While((fun () -> e.MoveNext()), (fun res -> (body e.Current) res))))

    member inline this.TryRun([<InlineIfLambda>] f: Code<_>) =
        let res = ref NoValue
        f res
        res.Value.Close()
    member inline this.Run ([<InlineIfLambda>] f) =
        match this.TryRun f with
        | ValueSome x -> x
        | ValueNone -> failwith "No return value exists"
    member inline this.Run ([<InlineIfLambda>] f: Code<unit>) =
        this.TryRun f |> ignore

let block = BlockBuilder()

let Break: Code<'T> = (fun res -> res.Value <- Break)
let Continue: Code<'T> = (fun res -> res.Value <- Continue)


let countdownFrom k n = block {
    let mutable n = n
    while n > 0 do
        if n * n = k then return n // newbie will see a ton of example algorithms in programming books that need a short-circuiting return statement
        printfn "%d" n
        n <- n - 1
    return n
}

let ex01() =
    countdownFrom 49 10 |> printfn "returned: %A"

let ex02() = block {
    for x in 1 .. 10 do
        if (x % 3 = 0) then yield! Continue
        printfn "number = %d" x
}

let ex03() = block {
    let mutable x = 1
    while true do
      if (x % 4 = 0) then yield! Break
      printfn "number = %d" x
      x <- x + 1
}