add ANF IR and conversion pass

This commit is contained in:
tali 2023-11-29 13:44:41 -05:00
parent 3cd480fbee
commit 0e1562984b
3 changed files with 307 additions and 1 deletions

View File

@ -4,5 +4,6 @@ let () =
try
let syn = Spice.parse "val x = 3 val y = x + 1" in
Fmt.pr "%a\n" Spice.Syn.pp_modl syn
let ir = Spice.compile syn in
Fmt.pr "%a\n" Spice.Ir.pp_entrypoint ir
with Spice.Error msg -> Logs.err (fun m -> m "%s" msg)

172
lib/anf_pass.ml Normal file
View File

@ -0,0 +1,172 @@
module Env = Map.Make (String)
type binding =
| Var_arg of Ir.arg
| Var_ele of Ir.arg * Ir.name
| Var_fun of Ir.arg * Ir.name
let ( @@@ ) f g x = f (g x)
let fcf_mthd_name = "call"
let anf modl =
let next_uid = ref 0 in
let gen_id name =
let uid = !next_uid in
incr next_uid;
Ir.Id.{ name; uid }
in
let lift_rhs name rhs =
let tmp = gen_id name in
let ctx rest = Ir.Let (tmp, rhs, rest) in
Ir.Arg tmp, ctx
in
let rec lower_exp env = function
| Syn.Literal l -> lift_rhs "lit" (Ir.Literal l)
| Syn.Path p -> lower_path env p
| Syn.Call (fn, args) ->
let (ob_v, ob_ctx), mthd =
match fn with
| Syn.Ele (ob, mthd) -> lower_exp env ob, mthd
| Syn.Var f -> (
match Env.find f env with
| Var_fun (ob, mthd) -> (ob, Fun.id), mthd
| _ -> lower_path env fn, fcf_mthd_name)
in
let arg_vs_rev, ctx =
List.fold_left
(fun (vs, ctx) arg ->
let v, v_ctx = lower_exp env arg in
v :: vs, ctx @@@ v_ctx)
([], ob_ctx)
args
in
let fn_e = Ir.Ele (ob_v, mthd) in
let arg_vs = List.rev arg_vs_rev in
let v, call_ctx = lift_rhs mthd (Ir.Call (fn_e, arg_vs)) in
v, ctx @@@ call_ctx
| Syn.Binop (op, e1, e2) ->
let v1, v1_ctx = lower_exp env e1 in
let v2, v2_ctx = lower_exp env e2 in
let v, op_ctx = lift_rhs "op" (Ir.Binop (op, v1, v2)) in
v, v1_ctx @@@ v2_ctx @@@ op_ctx
| Syn.If (e1, e2, e3) ->
let v1, v1_ctx = lower_exp env e1 in
let v2, v2_ctx = lower_exp env e2 in
let v3, v3_ctx = lower_exp env e3 in
let jn = gen_id "jn" in
let rv = gen_id "rv" in
let ctx rest =
let cont = Ir.{ params = [ rv ]; body = rest } in
let e2 = v2_ctx (Ir.Jump (jn, [ v2 ])) in
let e3 = v3_ctx (Ir.Jump (jn, [ v3 ])) in
v1_ctx (Ir.Let (jn, Cont cont, If (v1, e2, e3)))
in
Ir.Arg rv, ctx
| Syn.Obj items -> lower_obj env items
(* | Syn.Fun (xs, e) -> () *)
(* | Syn.Scope items -> () *)
| _ -> failwith "..."
and lower_path env = function
| Syn.Var x -> (
match Env.find x env with
| Var_arg v -> v, Fun.id
| Var_ele (ob, el) -> lift_rhs x (Ir.Get (Ele (ob, el)))
| Var_fun (ob, el) ->
(* TODO: special treatment for known FCF's? *)
lift_rhs x (Ir.Get (Ele (ob, el)))
| exception Not_found -> failwith (Fmt.str "undefined variable %S" x))
| Syn.Ele (ob, el) ->
let ob_v, ob_ctx = lower_exp env ob in
let el_v, el_ctx = lift_rhs el (Ir.Get (Ele (ob_v, el))) in
el_v, ob_ctx @@@ el_ctx
and lower_obj env items =
(* TODO: detect duplicate item names *)
let ob = gen_id "ob" in
let ob_v = Ir.Arg ob in
(* TODO: get last expression in block *)
let env, slots, mthds, inits =
List.fold_left
(fun (env, slots, mthds, inits) -> function
| Syn.Item_exp e ->
let inits = `Exp e :: inits in
env, slots, mthds, inits
| Syn.Item_val (name, rhs) ->
let env = Env.add name (Var_ele (ob_v, name)) env in
let slots = name :: slots in
let inits = `Val (name, rhs) :: inits in
env, slots, mthds, inits
| Syn.Item_obj (name, body) ->
let env = Env.add name (Var_ele (ob_v, name)) env in
let slots = name :: slots in
let inits = `Obj (name, body) :: inits in
env, slots, mthds, inits
| Syn.Item_fun (name, param_names, body) ->
let env = Env.add name (Var_fun (ob_v, name)) env in
let mthds = (name, param_names, body) :: mthds in
env, slots, mthds, inits)
(env, [], [], [])
items
in
(* mthds was built in reverse order, so rev_map fixes the order *)
let mthds =
List.rev_map
(fun (name, param_names, body) ->
(* TODO: detect duplicate param names *)
let params = List.map gen_id param_names in
let env' =
List.fold_left2
(fun env name id -> Env.add name (Var_arg (Arg id)) env)
env
param_names
params
in
let body_v, body_ctx = lower_exp env' body in
let body = body_ctx (Ir.Ret body_v) in
Ir.{ name; defn = { params; body } })
mthds
in
(* inits was built in reverse order, so make sure to *prepend* contexts
* when building the final context *)
let init_ctx =
List.fold_left
(fun ctx' -> function
| `Exp e ->
let _, ctx = lower_exp env e in
ctx @@@ ctx'
| `Val (name, rhs) ->
let rhs_v, rhs_ctx = lower_exp env rhs in
let set_ctx rest = Ir.Set (Ele (ob_v, name), rhs_v, rest) in
rhs_ctx @@@ set_ctx @@@ ctx'
| `Obj (name, body) ->
let ob_v, ob_ctx = lower_obj env body in
let set_ctx rest = Ir.Set (Ele (ob_v, name), ob_v, rest) in
ob_ctx @@@ set_ctx @@@ ctx')
Fun.id
inits
in
let slots = List.rev slots in
let ob_ctx rest = Ir.Let (ob, Obj { mthds; slots }, rest) in
ob_v, ob_ctx @@@ init_ctx
in
let std = gen_id "std" in
let std_env = Env.empty in
let std_env =
List.fold_left
(fun env fn -> Env.add fn (Var_fun (Arg std, fn)) env)
std_env
[ "read"; "print"; "itoa" ]
in
let std_env =
List.fold_left
(fun env var -> Env.add var (Var_ele (Arg std, var)) env)
std_env
[ "fs"; "rand" ]
in
let ret_v, ctx = lower_obj std_env modl.Syn.items in
let params = [ std ] in
let body = ctx (Ir.Ret ret_v) in
Ir.{ params; body }

133
lib/ir.ml Normal file
View File

@ -0,0 +1,133 @@
module Id = struct
type t = {
name : string;
uid : int;
}
let to_string { name; uid } = Fmt.str "%s#%d" name uid
let compare a b = Int.compare a.uid b.uid
let hash a = a.uid
let equal a b =
if a.uid = b.uid then (
assert (String.equal a.name b.name);
true)
else false
end
type name = Syn.name
type literal = Syn.literal
type binop = Syn.binop
(* TODO: split between arithmetic operators vs comparison operators *)
type arg = Arg of Id.t [@@unboxed]
type ele = Ele of arg * name
type exp =
| Let of Id.t * rhs * exp
| Set of ele * arg * exp
| If of arg * exp * exp
| Jump of Id.t * arg list
| Ret of arg
and rhs =
| Literal of literal
| Get of ele
| Call of ele * arg list
| Binop of binop * arg * arg
| Cont of cont
| Obj of block
and block = {
slots : name list;
mthds : mthd list;
}
and mthd = {
name : string;
defn : cont;
}
and cont = {
params : Id.t list;
body : exp;
}
type entrypoint = cont
(* pretty printer *)
let pp_list pp_ele ppf list =
Fmt.pf ppf "[";
List.iteri
(fun i ele ->
if i > 0 then Fmt.pf ppf ",";
pp_ele ppf ele)
list;
Fmt.pf ppf "]"
let pp_name ppf name = Fmt.pf ppf "%S" name
let rec pp_exp ppf exp =
Fmt.pf ppf "[";
let rec loop i exp =
if i > 0 then Fmt.pf ppf ",";
match exp with
| Let (name, rhs, rest) ->
Fmt.pf ppf "{\"let\":%a%a}" pp_id name pp_rhs rhs;
loop (i + 1) rest
| Set (e, v, rest) ->
Fmt.pf ppf "{\"set\":%a,\"value\":%a}" pp_ele e pp_arg v;
loop (i + 1) rest
| If (v1, e2, e3) ->
Fmt.pf ppf "{\"if\":%a,\"then\":%a,\"else\":%a}]" pp_arg v1 pp_exp e2 pp_exp e3
| Jump (tgt, args) ->
Fmt.pf ppf "{\"jump\":%a,\"args\":%a}]" pp_id tgt (pp_list pp_arg) args
| Ret v -> Fmt.pf ppf "{\"ret\":%a}]" pp_arg v
in
loop 0 exp
and pp_id ppf id = Fmt.pf ppf "%S" (Id.to_string id)
and pp_arg ppf = function
| Arg id -> pp_id ppf id
and pp_ele ppf = function
| Ele (Arg ob, el) -> Fmt.pf ppf "\"%s.%s\"" (Id.to_string ob) el
and pp_rhs ppf = function
| Literal (Int n) -> Fmt.pf ppf ",\"int\":%s" (Int64.to_string n)
| Literal l -> Fmt.pf ppf ",\"lit\":%S" (Syn.string_of_literal l)
| Get e -> Fmt.pf ppf ",\"get\":%a" pp_ele e
| Call (fn, args) ->
Fmt.pf ppf ",\"call\":%a,\"args\":%a" pp_ele fn (pp_list pp_arg) args
| Binop (op, v1, v2) ->
Fmt.pf
ppf
",\"binop\":%S,\"lhs\":%a,\"rhs\":%a"
(Syn.string_of_binop op)
pp_arg
v1
pp_arg
v2
| Cont { params; body } ->
Fmt.pf ppf ",\"cont\":%a,\"body\":%a" (pp_list pp_id) params pp_exp body
| Obj { slots; mthds } ->
Fmt.pf ppf ",\"obj\":{\"slots\":%a,\"mthds\":{" (pp_list pp_name) slots;
List.iteri
(fun i { name; defn = { params; body } } ->
if i > 0 then Fmt.pf ppf ",";
Fmt.pf
ppf
"%S:{\"params\":%a,\"body\":%a}"
name
(pp_list pp_id)
params
pp_exp
body)
mthds;
Fmt.pf ppf "}}"
let pp_entrypoint ppf { params; body } =
Fmt.pf ppf "{\"inputs\":%a,\"program\":%a}" (pp_list pp_id) params pp_exp body