small improvements to compilation

* don't eval the root of the updated graph when building a
supercombinator
* fix type-checking order resolution

compile.ml

@@ -194,13 +194,24 @@ and compile_strict (env : M.t string slot) = function
   | App (App (Ref f, x), y) as e ->
       match M.lookup f !prim_scs with
       | Some op when is_arith_op op ->
-          print ("compiling", f, "specially")
           compile_strict env x
             # compile_strict (incr <$> env) y
             # (op ::)
       | _ -> compile_lazy env e # (Eval ::)
   | e -> compile_lazy env e # (Eval ::)
 
+and compile_tail (env : M.t string slot) = function
+  | Ref v ->
+    match M.lookup v env with
+    | Some (As i) -> (Push (Arg i) ::)
+    | Some (Ls i) -> (Push (Local i) ::)
+    | None -> (Push (Combinator v) ::)
+  | App (f, x) ->
+    let f = compile_tail env f
+    let x = compile_lazy (map incr env) x
+    f # x # (Mkap ::)
+  | e -> compile_strict env e
+
 and compile_let cont env vs e =
   let n = length vs
   let env =
@@ -218,7 +229,7 @@ and compile_let cont env vs e =
 
 let supercomb (_, args, exp) =
   let env = M.from_list (zip args [0..length args])
-  let k = compile_strict (M.from_list (zip args (As <$> [0..length args]))) exp
+  let k = compile_tail (M.from_list (zip args (As <$> [0..length args]))) exp
   k [Update (length env), Pop (length env), Unwind]
 
 let compile_cons =
@@ -267,7 +278,6 @@ let program decs =
       | Foreign (Fimport { cc = Prim, var } as fi) ->
         define var (
           let (code, h) = cg_prim fi
-          print h
           prim_scs := M.insert var h !prim_scs
           pure [code]
         )

driver.ml

@@ -19,7 +19,8 @@ let go infile outfile =
     match lex prog str with
     | Right (ds, _) ->
       ds
-        |> T.tc_program
+        |> T.tc_program [] []
+        |> fun (_, _, z) -> z
         |> C.program
         |> A.assm_program
         |> write_bytes outfile
@@ -35,7 +36,11 @@ let go' infile outfile =
 let test str =
   match lex prog str with
   | Right (ds, _) ->
-      let code = ds |> T.tc_program |> C.program
+      let code =
+        ds
+          |> T.tc_program [] []
+          |> fun (_, _, z) -> z
+          |> C.program
       let lua = code |> A.assm_program
       print code
       put_line lua

lang.ml

@@ -7,6 +7,7 @@ type expr =
   | App  of expr * expr
   | Lam  of string * expr
   | Case of expr * list (string * list string * expr)
+  | If   of expr * expr * expr
   | Lit  of int
   | Let  of list (string * expr) * expr
 
@@ -29,6 +30,8 @@ let rec free_vars = function
         |> foldl S.union S.empty
         |> S.union (free_vars b)
         |> flip S.difference bound
+  | If (f, x, y) ->
+      S.union (free_vars f) (S.union (free_vars x) (free_vars y))
   | Lit _ -> S.empty
 
 type hstype =

lib/graph.ml

@@ -61,6 +61,13 @@ let toposort (graph : t 'a) : list 'a =
   iter visit nodes
   reverse !l
 
+let dot_of_graph (graph : t 'a) =
+  let mk node =
+    S.foldr (fun edge r -> show node ^ " -> " ^ show edge ^ "\n" ^ r) "\n"
+  "strict digraph {"
+    ^ M.foldr_with_key (fun node edges r -> mk node edges ^ r) "" graph
+    ^ "}"
+
 let groups_of_sccs (graph : t 'a) =
   let sccs = sccs graph
   let edges_of n =
@@ -72,10 +79,12 @@ let groups_of_sccs (graph : t 'a) =
       |> M.assocs
       |> map (fun (k, s) -> M.singleton s (S.singleton k))
       |> foldl (M.union_by (fun _ -> S.union)) M.empty
+  let atd nodes =
+    S.foldr (fun n -> S.union (edges_of n)) S.empty nodes `S.difference` nodes
   let comp_deps =
      components
        |> M.assocs
-       |> map (fun (node, edges) -> (node, edges_of node `S.difference` edges))
+       |> map (fun (node, edges) -> (node, atd edges))
        |> M.from_list
   let ordering = toposort comp_deps
   [ x | with k <- ordering, with Some x <- [M.lookup k components] ]

tc.ml

@@ -383,6 +383,7 @@ let dependency_graph defs =
           |> M.insert name S.empty
           |> (, define name x defs)
   let (graph, defs) = foldl go (M.empty, M.empty) defs
+  G.dot_of_graph graph |> put_line
   (G.groups_of_sccs graph, defs)
 
 let mk_lam args body = foldr (curry Lam) body args
@@ -410,17 +411,16 @@ let rec add_missing_vars scope = function
   | Tyarr (a, b) -> add_missing_vars (add_missing_vars scope b) a
   | Tytup xs -> foldl add_missing_vars scope xs
 
-let tc_program (prog : list decl) =
+let tc_program value_exports type_exports (prog : list decl) =
   let (plan, defs) = dependency_graph prog
+
   let tc_one (dt_info, val_scope, ty_scope, out) group =
-    print (length out, length (S.members group))
     let defs = [ x | with name <- S.members group, with Some x <- [M.lookup name defs] ]
     match defs with
     | [] -> (dt_info, val_scope, ty_scope, out)
     | [Foreign (Fimport {var, ftype}) as def] ->
         let ty_scope' = add_missing_vars M.empty ftype
         let t = check_is_type (M.union ty_scope' ty_scope) ftype
-        print var
         (dt_info, M.insert var (Forall { vars = M.keys ty_scope', body = t } |> Poly) val_scope, ty_scope, def :: out)
     | Cons (Foreign (Fimport {var}), _) ->
         error @@ "Foreign definition " ^ var ^ " is part of a group.  How?"
@@ -431,13 +431,11 @@ let tc_program (prog : list decl) =
         let (bindings, scope') = infer_binding_group dt_info -1 val_scope bindings
         let decs =
           [ Decl (name, unlambda expr) | with (name, expr) <- bindings ]
-        print name
         (dt_info, M.union (map force scope') val_scope, ty_scope, foldr (::) decs out)
     | Cons (Data d, ds) ->
         let datas = d :: [ d | with Data d <- ds ]
         let r = infer_data_group_kind ty_scope datas
         let fix_constr (name, rhos : list tc_rho) =
-          print name
           Constr (name, replicate (length rhos) (Tycon "#"))
         let rec go dt ty (vl : M.t string (scheme tc_rho)) ds = function
           | [] -> (dt, vl, ty, ds)
@@ -452,5 +450,9 @@ let tc_program (prog : list decl) =
                 (Data (name, [], fix_constr <$> constrs) :: ds)
                 rest
         go dt_info ty_scope val_scope out r
-  let (_, _, _, p) = foldl tc_one (M.empty, M.empty, M.empty, []) plan
-  p
+  let (_, vals, types, p) = foldl tc_one (M.empty, M.empty, M.empty, []) plan
+  (
+    [x | with k <- value_exports, with Some x <- [M.lookup k vals]],
+    [x | with k <- type_exports, with Some x <- [M.lookup k types]],
+    p
+  )