Add structures.

1 files changed, 76 insertions, 17 deletions

diff --git a/elab.sml b/elab.sml
index a99758a..33f1a10 100644
--- a/elab.sml
+++ b/elab.sml
@@ -35,34 +35,47 @@ struct
         NONE => []
       | SOME (heads, tails) => heads :: transpose tails
 
-  type env = int StringMap.map * (int * int) StringMap.map
+  datatype env = Env of { vars: int StringMap.map, types: (int * int) StringMap.map, structTypes: env StringMap.map }
 
-  fun bindVar (name : string) (sym : int) ((vars, types) : env) : env = (StringMap.insert name sym vars, types)
+  val emptyEnv = Env { vars = StringMap.empty, types = StringMap.empty, structTypes = StringMap.empty }
 
-  fun lookupVar (name : string) ((vars, _) : env) : int =
-    case StringMap.lookup name vars of
+  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) ((vars, types) : env) : env =
+  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, types)
+        (0, #vars env, #types env)
         cons
-  in (vars, types)
+  in Env { vars = vars, types = types, structTypes = #structTypes env }
   end
 
-  fun lookupCon (name : string) ((_, types) : env) : int option =
-    Option.map (fn (i, _) => i) (StringMap.lookup name types)
+  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) ((_, types) : env) : int =
-    let val (_, n) = valOf (StringMap.lookup name types)
-    in n
-    end
+  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 =
@@ -179,8 +192,8 @@ struct
     case expr of
       Syntax.PTuple t => Syntax.PTuple (map (fixConstructors env) t)
     | Syntax.PVar v =>
-        if isSome (lookupCon v env)
-        then Syntax.PCon (v, Syntax.PTuple [])
+        if isSome (lookupCon [v] env)
+        then Syntax.PCon ([v], Syntax.PTuple [])
         else Syntax.PVar v
     | Syntax.PCon (s, p) => Syntax.PCon (s, fixConstructors env p)
     | _ => expr
@@ -236,6 +249,32 @@ struct
               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) =>
@@ -280,7 +319,18 @@ struct
   and elab (env : env) (p : Syntax.expr) : Syntax.lexp =
     case p of
       Syntax.EIdent [i] => Syntax.LVar (lookupVar i env)
-    | Syntax.EIdent _ => raise Fail "long identifiers are not supported"
+    | 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.EBuiltin builtin => Syntax.LPrim (primop builtin)
     | Syntax.EInt i => Syntax.LInt i
     | Syntax.EStr s => Syntax.LString s
@@ -359,6 +409,15 @@ struct
             )
         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])
@@ -368,5 +427,5 @@ struct
         in Syntax.LApp (Syntax.LFn (v, elabCase env (Syntax.LVar v) arms), elab env expr)
         end
 
-  fun elaborate (p : Syntax.expr) : Syntax.lexp = elab (StringMap.empty, StringMap.empty) p
+  fun elaborate (p : Syntax.expr) : Syntax.lexp = elab emptyEnv p
 end