A simple garbage-collected language using mark-and-sweep

gc.ml

(**********************************************************************************)
(* micro-c *)

type microc =
  | Let of { id: string; binding: microc; body: microc }
  | Var of string
  | Malloc
  | Gc
  | Free of microc
  | Set of {loc: microc; value: microc}
  | Deref of {loc: microc}
  | Num of int

type value =
    VNum of int
  | VLoc of int

type free_cell = Free | Occupied

type heap = {
  (* array of indices that are occupied in the heap; true = not-occupied *)
  free_list: free_cell array;
  heap: value array
}

let heap_size = 100


let empty_heap () = { free_list = Array.make heap_size Free; heap = Array.make heap_size (VNum (-1))}

module StringMap = Map.Make(String)

type microcenv = value StringMap.t

(* out of memory exception *)
exception Oom

(* out of memory exception *)
exception Runtime

(* find a free index in the free_list or raise an Oom exception *)
let find_free_idx (h:heap) : int =
  let rec helper (h:heap) (idx:int) : int =
    if (idx >= Array.length h.free_list) then raise Oom else
      if (Array.get h.free_list idx) = Free then idx else helper h (idx+1) in
  helper h 0

let rec gc (root:value) (heap:heap) : unit =
  (* traverse the heap beginning from `root` and mark every reachable address as occupied *)
  match root with
  | VNum(_) -> ()
  | VLoc(l) ->
    (* check to see if we've already marked this as free; if not, recurse *)
    if (Array.get heap.free_list l) = Free then
      (Array.set heap.free_list l Occupied;
       gc (Array.get heap.heap l) heap)

let rec interp_c (heap:heap) (env: microcenv) (microc:microc) : value =
  match microc with
  | Let {id; binding; body} ->
    let bindv = interp_c heap env binding in
    let new_env = StringMap.add id bindv env in
    interp_c heap new_env body
  | Var(s) ->
    (match StringMap.find_opt s env with
     | Some(v) -> v
     | None -> raise Runtime)
  | Malloc ->
    let idx = find_free_idx heap in
    (* update free list *)
    Array.set heap.free_list idx Occupied;
    VLoc(idx)
  | Set {loc; value} ->
    let l = (match interp_c heap env loc with
          VLoc(v) -> v
        | VNum(_) -> raise Runtime) in
    let v = interp_c heap env value in
    Array.set heap.heap l v;
    VNum(0)
  | Deref { loc } ->
    let l = (match interp_c heap env loc with
       VLoc(v) -> v
     | VNum(_) -> raise Runtime) in
    Array.get heap.heap l
  | Num(n) -> (VNum(n))
  | Gc ->
    (* mark all heap locations as free *)
    Array.iteri (fun idx _ -> Array.set heap.free_list idx Free) heap.free_list;
    (* then, for every location reachable in our environment, mark it as occupied *)
    StringMap.iter (fun _ value -> gc value heap) env;
    VNum(0)
  | Free(exp) ->
    let l = (match interp_c heap env exp with
          VLoc(v) -> v
        | VNum(_) -> raise Runtime) in
    Array.set heap.free_list l Free;
    VNum(0)


(**********************************************************************************)

exception Type_error

type ast =
    Let of  string * ast * ast
  | Unbox of ast
  | Box of ast
  | Set of string * ast
  | Var of string
  | Num of int

let fresh_name f =
  let v = Format.sprintf "fresh%d" (!f) in
  f := !f + 1;
  v

(* compile ast into microc *)
let rec compile_leaky (ast:ast) (f:int ref) : microc =
  match ast with
  | Box(e) ->
    let e_v = compile_leaky e f in
    (* whenever we allocate a new value, garbage collect *)
    let loc_name = fresh_name f in
    Let{id=loc_name; binding=Malloc;
       body = Let {id="_"; binding=Set { loc=Var(loc_name); value=e_v };
          body=Var(loc_name) }}
  | Num(n) -> Num(n)
  | Var(s) -> Var(s)
  | Let(id, binding, body) ->
    Let {id=id; binding=compile_leaky binding f; body= compile_leaky body f}
  | Set(loc, value) ->
    Set { loc = Var(loc); value = compile_leaky value f }
  | Unbox(e) ->
    Deref { loc = compile_leaky e f }

let compile_and_run_leaky ast =
  let compiled = compile_leaky ast (ref 0) in
  interp_c (empty_heap ()) StringMap.empty compiled

let p1 =
  Let("v", Num(10),
      Let("x", Box(Var("v")),
          Let("_", Set("x", Num(20)),
              Unbox(Var("x")))))

let p2 = Let("x", Box(Num(10)), Unbox(Var("x")))

let () =
  assert ((compile_and_run_leaky p1) = VNum(20))

(* compile with gc *)
let rec compile_gc (ast:ast) (f:int ref) : microc =
  match ast with
  | Box(e) ->
    (* whenever we allocate a new value, garbage collect *)
    let e_v = compile_leaky e f in
    let loc_name = fresh_name f in
    let res : microc = Let{id=loc_name; binding=Malloc;
                           body = Let {id="_"; binding=Set { loc=Var(loc_name); value=e_v };
                                       body=Var(loc_name) }} in
    let body : microc = Let { id = "_"; binding = Gc; body = Var("res")} in
    Let { id = "res"; binding = res; body=body }
  | Num(n) -> Num(n)
  | Var(s) -> Var(s)
  | Let(id, binding, body) ->
    Let {id=id; binding=compile_gc binding f; body= compile_gc body f}
  | Set(loc, value) ->
    Set { loc = Var(loc); value = compile_gc value f }
  | Unbox(e) ->
    Deref { loc = compile_gc e f }