A small STG-like lazy abstract machine

stg.ml

(** A small STG-like lazy abstract machine. What I take from the STG
    is that constructors are responsible for choosing a branch in a
    pattern-matching (the T in STG) and for updating themselves (by
    pushing an update frame). Apart from that it's not very different
    from a simple KAM. Note that the current implementation in GHC is
    not actually tagless anymore, as it proved to be slower than a
    tagged mechanism. But since it's a simpler design, (and the
    difference doesn't matter much unless actual machine code is
    emitted).

    To capture the essence of the abstract machine I use a λ-calculus
    with pairs and sums. I don't use a dedicated subset with [let]
    forms expressing sharing. Like many of the details of the actual
    STG it is most relevant when emitting machine code.

    To represent sharing (a.k.a. the heap) I use Ocaml's heap and
    update references.

    Reference: the STG paper:
    - https://www.cambridge.org/core/journals/journal-of-functional-programming/article/div-classtitleimplementing-lazy-functional-languages-on-stock-hardware-the-spineless-tagless-g-machinediv/354FFB29102309CCD2A3824F894A2799
    - http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.53.3729&rep=rep1&type=pdf

*)

(* Untested yet. I've just written it down top to bottom. *)


module Id = struct
  type t = string

  module Map = Map.Make(struct type t = string let compare = compare end)
end

(** [term] are terms of the source language. It will also serve as
    code for the abstract machine. *)
(* TODO: add fixpoints? *)
type term =
  | Var of Id.t
  | Lam of Id.t * term
  | App of term * term
  | Pair of term * term
  | CasePair of term * Id.t * Id.t * term
  | Iota1 of term
  | Iota2 of term
  | CaseSum of term * Id.t * term * Id.t * term

(** [closure] is the type of values manipulated by the abstract
    machine. It is a reference because it may be updatable. The
    simplifying idea of the STG is that the instructions do not
    distinguish between value which are updatable or not. Only the
    values themselves do.

    Another note is that I chose to have continuations frames as a
    kind of values. There are other mostly equivalent design choices
    (having a different stack for continuation frames, or having the
    stack itself distinguish between continuation frames and actual
    values). Continuation frames are either a pattern matching or an
    update frame. *)
type closure =
  value ref
and value =
  | Fun of env * Id.t * term
  | Thunk of env * term
  | TPair of closure * closure
  | TIota1 of closure
  | TIota2 of closure
  | Cont of cont
and env =
  closure Id.Map.t
and cont =
  | KPair of env * Id.t * Id.t * term
  | KSum  of env * Id.t * term * Id.t * term
  | Update of closure

type stack = closure list

type state = {
  env : env;
  code : term;
  stack : stack;
}

let alloc env u = ref @@ Thunk(env,u)

let rec step : state -> state = fun {env;code;stack} ->
  match code with
  | Var id ->
    continue (Id.Map.find id env) stack
  | Lam (x,u) ->
    continue (ref @@ Fun (env,x,u)) stack
  | App (u,v) ->
    { env ; code=u ; stack = (alloc env u)::stack }
  | Pair (u,v) ->
    continue (ref @@ TPair (alloc env u, alloc env v)) stack
  | CasePair (u,x,y,k) ->
    { env ; code = u ; stack = (ref @@ Cont (KPair (env,x,y,k)))::stack }
  | Iota1 u ->
    continue (ref @@ TIota1 (alloc env u)) stack
  | Iota2 v ->
    continue (ref @@ TIota2 (alloc env v)) stack
  | CaseSum (u,x,k1,y,k2) ->
    { env ; code = u ; stack = (ref @@ Cont (KSum (env,x,k1,y,k2)))::stack }
and continue : closure -> stack -> state = fun cl stack ->
  match !cl , stack with
  | v , { contents = Cont (Update cl') }::rest ->
    cl' := v ; continue cl rest
  | Fun (env,x,code) , arg::rest ->
    { env = Id.Map.add x arg env ; code ; stack = rest }
  | Thunk (env,code) , stack ->
    (* To implement black-holing, we could update [cl] here with a
       black-hole value. *)
    { env ; code ; stack = (ref @@ Cont (Update cl))::stack }
  | TPair (cl1,cl2) , { contents = Cont (KPair (env,x,y,code)) }::rest ->
    let ext = env |> Id.Map.add x cl1 |> Id.Map.add y cl2 in
    { env = ext ; code ; stack = rest }
  | TIota1 (cl) , { contents = Cont (KSum (env,x,code,_,_)) }::rest ->
    { env = Id.Map.add x cl env ; code ; stack = rest }
  | TIota2 (cl) , { contents = Cont (KSum (env,_,_,y,code)) }::rest ->
    { env = Id.Map.add y cl env ; code ; stack = rest }
  | _ -> failwith "Typing error"