diff options
Diffstat (limited to 'Elab.sml')
| -rw-r--r-- | Elab.sml | 419 |
1 files changed, 419 insertions, 0 deletions
diff --git a/Elab.sml b/Elab.sml new file mode 100644 index 0000000..60760cd --- /dev/null +++ b/Elab.sml @@ -0,0 +1,419 @@ +structure Elab = +struct + structure StringMap = Map(type k = string val cmp = String.compare) + structure IntMap = Map(type k = int val cmp = Int.compare) + + fun primop (s : string) : Syntax.primop = + case s of + "exit" => Syntax.PExit + | "add" => Syntax.PAdd + | "sub" => Syntax.PSub + | "mul" => Syntax.PMul + | "div" => Syntax.PDiv + | "read" => Syntax.PRead + | "write" => Syntax.PWrite + | "writeErr" => Syntax.PWriteErr + | _ => raise Fail ("invalid op: " ^ s) + + fun enumerate (l : 'a list) : (int * 'a) list = + let + fun go _ [] = [] + | go i (x :: xs) = (i, x) :: go (i + 1) xs + in go 0 l + end + + fun hdstls ([] : 'a list list) : ('a list * 'a list list) option = SOME ([], []) + | hdstls ([] :: _) = NONE + | hdstls ((x :: xs) :: ls) = + case hdstls ls of + NONE => NONE + | SOME (heads, tails) => SOME (x :: heads, xs :: tails) + + fun transpose ([] : 'a list list) : 'a list list = [] + | transpose l = + case hdstls l of + NONE => [] + | SOME (heads, tails) => heads :: transpose tails + + datatype env = Env of { vars: int StringMap.map, types: (int * int) StringMap.map, structTypes: env StringMap.map } + + val emptyEnv = Env { vars = StringMap.empty, types = StringMap.empty, structTypes = StringMap.empty } + + fun bindVar (name : string) (sym : int) (Env env) : env = Env { vars = StringMap.insert name sym (#vars env), types = #types env, structTypes = #structTypes env } + + fun lookupVar (name : string) (Env env) : int = + case StringMap.lookup name (#vars env) of + SOME x => x + | NONE => raise Fail ("unbound identifier " ^ name) + + fun bindDataCons (cons : (string * Syntax.etype option) list) (Env env) : env = + let + val nCons = length cons + val (_, vars, types) = + foldl + (fn ((name, _), (i, vars, types)) => + (i + 1, StringMap.insert name (GenSym.new ()) vars, StringMap.insert name (i, nCons) types)) + (0, #vars env, #types env) + cons + in Env { vars = vars, types = types, structTypes = #structTypes env } + end + + fun lookupStructType (name : string) (Env env) : env = + case StringMap.lookup name (#structTypes env) of + SOME x => x + | NONE => raise Fail ("unbound structure " ^ name) + + fun lookupCon ([name] : string list) (Env env) : int option = + Option.map (fn (i, _) => i) (StringMap.lookup name (#types env)) + | lookupCon (structName :: names) env = + lookupCon names (lookupStructType structName env) + | lookupCon [] _ = raise Fail "lookupCon empty" + + fun nConstructors ([name] : string list) (Env env) : int = + let val (_, n) = valOf (StringMap.lookup name (#types env)) + in n + end + | nConstructors (structName :: names) env = + nConstructors names (lookupStructType structName env) + | nConstructors [] _ = raise Fail "nConstructors empty" + + fun patternMatrix (arms : Syntax.pat list) : Syntax.pat list list = + let val ts = + foldl + (fn (Syntax.PWild, tupleSize) => tupleSize + | (Syntax.PVar _, tupleSize) => tupleSize + | (Syntax.PCon _, _) => ~2 + | (Syntax.PInt _, _) => ~2 + | (Syntax.PTuple t, ~1) => length t + | (Syntax.PTuple t, tupleSize) => + if tupleSize = length t + then tupleSize + else raise Fail "a type checker would have caught this") + ~1 + arms + in if ts < 0 + then [arms] + else let + val expandedArms = + map + (fn Syntax.PTuple t => t + | _ => List.tabulate (ts, fn _ => Syntax.PWild)) + arms + val cols = transpose expandedArms + in List.concat (map patternMatrix cols) + end + end + + fun occurrenceVector (expr : Syntax.lexp) (arms : Syntax.pat list) : Syntax.lexp list = + let + val expandedArms = + List.mapPartial + (fn Syntax.PTuple t => SOME t + | _ => NONE) + arms + val cols = transpose expandedArms + in + case cols of + [] => [expr] + | _ => List.concat (map (fn (i, col) => occurrenceVector (Syntax.LSelect (i, expr)) col) (enumerate cols)) + end + + fun patternBindings (expr : Syntax.lexp) (Syntax.PVar v) : (string * Syntax.lexp) list = [(v, expr)] + | patternBindings expr (Syntax.PTuple t) = + List.concat (map (fn (i, p) => patternBindings (Syntax.LSelect (i, expr)) p) (enumerate t)) + | patternBindings expr (Syntax.PCon (_, arg)) = patternBindings (Syntax.LSelect (1, expr)) arg + | patternBindings _ _ = [] + + fun swap1 0 (l : 'a list) : 'a list = l + | swap1 n (first :: rest) = + (case swap1 (n - 1) rest of + x :: rest => x :: first :: rest + | _ => raise Fail "swap1: index out of bounds") + | swap1 _ _ = raise Fail "swap1: index out of bounds" + + fun swap (n : int) (patterns : Syntax.pat list list) (occurrences : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list = + (map (swap1 n) patterns, swap1 n occurrences) + + fun specialize (env : env) (n : int) (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list = + let + val con1 = + List.find + (fn Syntax.PCon (name, _) => valOf (lookupCon name env) = n + | _ => false) + (map hd patterns) + val newTupleSize = + case con1 of + SOME (Syntax.PCon (_, Syntax.PTuple t)) => length t + | _ => 0 + val occHead = hd occurrences + val occRest = tl occurrences + val occurrences = + if newTupleSize = 0 + then Syntax.LSelect (1, occHead) :: occRest + else + List.tabulate (newTupleSize, fn i => Syntax.LSelect (i, Syntax.LSelect (1, occHead))) @ occRest + fun specializeRow (Syntax.PInt i :: rest) = + if i = n then SOME (Syntax.PWild :: rest) else NONE + | specializeRow (Syntax.PWild :: rest) = SOME (Syntax.PWild :: rest) + | specializeRow (Syntax.PVar _ :: rest) = SOME (Syntax.PWild :: rest) + | specializeRow (Syntax.PCon (con, Syntax.PTuple []) :: rest) = + specializeRow (Syntax.PCon (con, Syntax.PTuple [Syntax.PWild]) :: rest) + | specializeRow (Syntax.PCon (con, Syntax.PTuple args) :: rest) = + if valOf (lookupCon con env) = n + then SOME (args @ rest) + else NONE + | specializeRow (Syntax.PCon (con, obj) :: rest) = + specializeRow (Syntax.PCon (con, Syntax.PTuple [obj]) :: rest) + | specializeRow _ = raise Fail "unexpected pattern in the matrix" + val (patterns, actions) = + ListPair.unzip + (List.mapPartial + (fn (p, a) => + Option.map (fn p => (p, a)) (specializeRow p)) + (ListPair.zipEq (patterns, actions))) + in (patterns, occurrences, actions) + end + + fun default (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list = + let + val (patterns, actions) = + ListPair.unzip + (List.filter + (fn (Syntax.PWild :: _, _) => true + | (Syntax.PVar _ :: _, _) => true + | _ => false) + (ListPair.zipEq (patterns, actions))) + in (patterns, occurrences, actions) + end + + fun compilePatternMatching (env : env) ([] : Syntax.pat list list, _ : Syntax.lexp list, _ : Syntax.lexp list) : Syntax.lexp = + raise Fail "nonexhaustive match" + | compilePatternMatching env (patterns as firstRow :: rows, occurrences, actions) = + let val refutablePattern = + List.find + (fn (_, Syntax.PInt _) => true + | (_, Syntax.PCon _) => true + | _ => false) + (enumerate firstRow) + in + case refutablePattern of + NONE => hd actions + | SOME (i, _) => + let + val (patterns, occurrences) = + if i = 0 + then (patterns, occurrences) + else swap i patterns occurrences + val firstCol = map hd patterns + val signatures = + map (fn (x, _) => x) + (IntMap.toList + (foldl + (fn (Syntax.PInt i, acc) => IntMap.insert i true acc + | (Syntax.PCon (c, _), acc) => IntMap.insert (valOf (lookupCon c env)) true acc + | (_, acc) => acc) + IntMap.empty + firstCol)) + val nCons = + case List.find (fn (Syntax.PCon _) => true | _ => false) firstCol of + SOME (Syntax.PCon (name, _)) => nConstructors name env + | _ => ~1 + val defaultCase = + if length signatures = nCons + then NONE + else SOME (compilePatternMatching env (default (patterns, occurrences, actions))) + val switchOperand = + if nCons < 0 + then hd occurrences + else Syntax.LSelect (0, hd occurrences) + in + Syntax.LSwitch + ( switchOperand + , map + (fn i => (i, compilePatternMatching env (specialize env i (patterns, occurrences, actions)))) + signatures + , defaultCase + ) + end + end + + fun declBoundVars (Syntax.DVal (p, _)) : string list = map (fn (x, _) => x) (patternBindings (Syntax.LInt 0) p) + | declBoundVars (Syntax.DValRec (p, _)) = map (fn (x, _) => x) (patternBindings (Syntax.LInt 0) p) + | declBoundVars (Syntax.DFun (name, _)) = [name] + | declBoundVars (Syntax.DDatatype (_, cons)) = map (fn (x, _) => x) cons + | declBoundVars (Syntax.DStruct (name, _)) = [name] + + fun structBoundVars (decls : Syntax.dec list) : string list = List.concatMap declBoundVars decls + + fun bindStructType (name : string) (decls : Syntax.dec list) (Env env) : env = + let + val structEnv = + foldl + (fn (Syntax.DDatatype (_, cons), env) => bindDataCons cons env + | (Syntax.DStruct (name, decls), env) => bindStructType name decls env + | (_, env) => env) + emptyEnv + decls + val structEnv = + foldl + (fn ((i, n), env) => bindVar n i env) + structEnv + (enumerate (structBoundVars decls)) + in + Env { vars = #vars env, types = #types env, structTypes = StringMap.insert name structEnv (#structTypes env) } + end + + fun actionVector (env : env) (expr : Syntax.lexp) (arms : (Syntax.pat * Syntax.expr) list) : Syntax.lexp list = + map + (fn (p, body) => + let + val bindings = patternBindings expr p + val env = + foldl + (fn ((name, _), env) => + bindVar name (GenSym.new ()) env) + env + bindings + in + foldl + (fn ((name, binding), acc) => + Syntax.LApp (Syntax.LFn (lookupVar name env, acc), binding)) + (elab env body) + bindings + end) + arms + + and elabCase (env : env) (expr : Syntax.lexp) (arms : (Syntax.pat * Syntax.expr) list) = + let + val patterns = transpose (patternMatrix (map (fn (x, _) => x) arms)) + val occurrences = occurrenceVector expr (map (fn (x, _) => x) arms) + val actions = actionVector env expr arms + val actionFns = + map + (fn a => (GenSym.new (), GenSym.new (), a)) + actions + val smallActions = + map + (fn (f, _, _) => Syntax.LApp (Syntax.LVar f, Syntax.LInt 0)) + actionFns + in + Syntax.LFix + ( actionFns + , compilePatternMatching env (patterns, occurrences, smallActions) + ) + end + + and elab (env : env) (p : Syntax.expr) : Syntax.lexp = + case p of + Syntax.EIdent [i] => Syntax.LVar (lookupVar i env) + | Syntax.EIdent (structName :: accessors) => + let + val s = Syntax.LVar (lookupVar structName env) + val env = lookupStructType structName env + fun go env [i] acc = Syntax.LSelect (lookupVar i env, acc) + | go env (accessor :: accessors) acc = + let val env = lookupStructType accessor env + in go env accessors (Syntax.LSelect (lookupVar accessor env, acc)) + end + | go _ [] _ = raise Fail "go empty" + in go env accessors s + end + | Syntax.EIdent [] => raise Fail "invalid syntax" + | Syntax.EBuiltin builtin => Syntax.LPrim (primop builtin) + | Syntax.EInt i => Syntax.LInt i + | Syntax.EStr s => Syntax.LString s + | Syntax.ETuple exprs => Syntax.LRecord (map (elab env) exprs) + | Syntax.EList exprs => + foldr + (fn (x, acc) => + Syntax.LRecord [elab env x, acc]) + (Syntax.LInt 0) + exprs + | Syntax.EApp (f, x) => Syntax.LApp (elab env f, elab env x) + | Syntax.ETyped (e, _) => elab env e + | Syntax.EAndAlso (_, _) => raise Fail "unimplemented" + | Syntax.EOrElse (_, _) => raise Fail "unimplemented" + | Syntax.ELet ([], body) => elab env body + | Syntax.ELet (Syntax.DDatatype (name, cons) :: decls, body) => + let + val env = bindDataCons cons env + val funs = + List.mapPartial + (fn (_, (_, NONE)) => NONE + | (i, (name, _)) => + let val v = GenSym.new () + in SOME (lookupVar name env, v, Syntax.LRecord [Syntax.LInt i, Syntax.LVar v]) + end) + (enumerate cons) + val vals = + List.mapPartial + (fn (i, (name, NONE)) => SOME (lookupVar name env, Syntax.LRecord [Syntax.LInt i]) + | _ => NONE) + (enumerate cons) + in + foldl + (fn ((v, x), acc) => + Syntax.LApp (Syntax.LFn (v, acc), x)) + (Syntax.LFix (funs, elab env (Syntax.ELet (decls, body)))) + vals + end + | Syntax.ELet (Syntax.DVal (pat, v) :: decls, body) => + elab env (Syntax.ECase (v, [(pat, Syntax.ELet (decls, body))])) + | Syntax.ELet (Syntax.DValRec (Syntax.PVar name, f as Syntax.ELambda _) :: decls, body) => + let + val n = GenSym.new () + val env = bindVar name n env + val (arg, fnBody) = + case elab env f of + Syntax.LFn x => x + | _ => raise Fail "Syntax.ELambda should expand to Syntax.LFn" + in + Syntax.LFix ([(n, arg, fnBody)], elab env (Syntax.ELet (decls, body))) + end + | Syntax.ELet (Syntax.DValRec _ :: _, _) => raise Fail "invalid val rec" + | Syntax.ELet (Syntax.DFun (name, cases) :: decls, body) => + let + val (ps1, _) = hd cases + val nPats = length ps1 + in if not (List.all (fn (ps, _) => length ps = nPats) cases) + then raise Fail "clauses do not all have same number of patterns" + else let + val n = GenSym.new () + val temps = List.tabulate (nPats, fn _ => GenSym.new ()) + val env = bindVar name n env + val t = GenSym.new () + val innerCase = elabCase env (Syntax.LVar t) (map (fn (ps, b) => (Syntax.PTuple ps, b)) cases) + in + Syntax.LFix + ( [ ( n + , hd temps + , foldr + Syntax.LFn + (Syntax.LApp (Syntax.LFn (t, innerCase), Syntax.LRecord (map Syntax.LVar temps))) + (tl temps) + ) + ] + , elab env (Syntax.ELet (decls, body)) + ) + end + end + | Syntax.ELet (Syntax.DStruct (name, structDecls) :: decls, body) => + let + val names = structBoundVars structDecls + val tuple = elab env (Syntax.ELet (structDecls, Syntax.ETuple (map (fn n => Syntax.EIdent [n]) names))) + val v = GenSym.new () + val env = bindStructType name structDecls env + val env = bindVar name v env + in Syntax.LApp (Syntax.LFn (v, elab env (Syntax.ELet (decls, body))), tuple) + end + | Syntax.ELambda body => + let val v = GenSym.new () + in Syntax.LFn (v, elabCase env (Syntax.LVar v) [body]) + end + | Syntax.ECase (expr, arms) => + let val v = GenSym.new () + in Syntax.LApp (Syntax.LFn (v, elabCase env (Syntax.LVar v) arms), elab env expr) + end + + fun elaborate (p : Syntax.expr) : Syntax.lexp = elab emptyEnv p +end |