don't lift cases in strict positions

compile.ml

@@ -1,5 +1,6 @@
 module M = import "data/map.ml"
 module S = import "data/set.ml"
+module Str = import "data/set.ml"
 
 open import "prelude.ml"
 open import "./lang.ml"
@@ -56,12 +57,12 @@ instance show program_item begin
     | Fd _ -> "<foreign item>"
 end
 
-let rec lambda_lift = function
+let rec lambda_lift strict = function
   | Ref v -> pure (Ref v)
   | Lit v -> pure (Lit v)
-  | App (f, x) -> (| app (lambda_lift f) (lambda_lift x) |)
+  | App (f, x) -> (| app (lambda_lift false f) (lambda_lift false x) |)
   | Lam (v, x) ->
-    let! body = lambda_lift x
+    let! body = lambda_lift true x
     let! (i, defs, known_sc) = get
 
     let vars =
@@ -76,23 +77,26 @@ let rec lambda_lift = function
     put (i + 1, Decl def :: defs, known_sc)
       |> map (const app)
   | Case (sc, alts) ->
-    let! sc = lambda_lift sc
-    let! alts = traverse (fun (c, args, e) -> (c,args,) <$> lambda_lift e) alts
+    let! sc = lambda_lift true sc
+    let! alts = traverse (fun (c, args, e) -> (c,args,) <$> lambda_lift true e) alts
     let case = Case (sc, alts)
-    let! (i, defs, known_sc) = get
-    let vars =
-      case
-        |> free_vars
-        |> flip S.difference known_sc
-        |> S.members
-    let def = ("Lam" ^ show i, vars, case)
-    let app = foldl (fun f -> app f # Ref) (Ref ("Lam" ^ show i)) vars
-    put (i + 1, Decl def :: defs, known_sc)
-      |> map (const app)
+    if strict then
+      pure case
+    else
+      let! (i, defs, known_sc) = get
+      let vars =
+        case
+          |> free_vars
+          |> flip S.difference known_sc
+          |> S.members
+      let def = ("Lam" ^ show i, vars, case)
+      let app = foldl (fun f -> app f # Ref) (Ref ("Lam" ^ show i)) vars
+      put (i + 1, Decl def :: defs, known_sc)
+        |> map (const app)
   | Let (vs, e) ->
     let! vs = flip traverse vs @@ fun (v, e) ->
-      (v,) <$> lambda_lift e
-    let! e = lambda_lift e
+      (v,) <$> lambda_lift false e
+    let! e = lambda_lift true e
     pure (Let (vs, e))
 
 let rec eta_contract = function
@@ -113,7 +117,7 @@ let rec lambda_lift_sc = function
     match e with
     | Lam (v, e) -> lambda_lift_sc (Decl (n, a ++ [v], e))
     | _ ->
-      let! e = lambda_lift e
+      let! e = lambda_lift true e
       let! _ = modify (fun (a, b, s) -> (a, b, S.insert n s))
       pure (Decl (n, a, e))
   | Data c -> Data c |> pure
@@ -130,14 +134,14 @@ let cg_prim (Fimport { var, fent }) =
     [ Push (Arg 0), Eval (* x, arg0, arg1, arg2, arg3 *)
     , Push (Arg 2), Eval (* y, x, arg0, arg1, arg2, arg3 *)
     , Sub                (* y - x, arg0, arg1, arg2, arg3 *)
-    , Iszero ([ Push (Arg 3) ], [ Push (Arg 4) ])
-    , Update 4, Pop 4, Unwind ]
+    , Iszero ([ Pack (0, 0) ], [ Pack (0, 1) ])
+    , Update 2, Pop 2, Unwind ]
   match fent with
   | "add" -> (Sc (var, 2, prim_math_op Add), Add)
   | "sub" -> (Sc (var, 2, prim_math_op Sub), Sub)
   | "mul" -> (Sc (var, 2, prim_math_op Mul), Mul)
   | "div" -> (Sc (var, 2, prim_math_op Div), Div)
-  | "equ" -> (Sc (var, 4, prim_equality), Unwind)
+  | "equ" -> (Sc (var, 2, prim_equality), Unwind)
   | "seq" -> (Sc (var, 2, [ Push (Arg 0), Eval, Update 0, Push (Arg 2), Update 2, Pop 2, Unwind]), Eval)
   | e -> error @@ "No such primitive " ^ e
 
@@ -146,9 +150,16 @@ type slot = As of int | Ls of int
 let offs n = function
   | As x -> As (x + n)
   | Ls x -> Ls (x + n)
+
 let incr = offs 1
 
-let rec compile (env : M.t string slot) = function
+let private prim_scs : ref (M.t string gm_code) = ref M.empty
+
+let private is_arith_op = function
+  | Add | Sub | Mul | Div | Iszero _ -> true
+  | _ -> false
+
+let rec compile_lazy (env : M.t string slot) = function
   | Ref v ->
     match M.lookup v env with
     | Some (As i) -> (Push (Arg i) ::)
@@ -156,12 +167,18 @@ let rec compile (env : M.t string slot) = function
     | None -> (Push (Combinator v) ::)
 
   | App (f, x) ->
-    let f = compile env f
-    let x = compile (map incr env) x
+    let f = compile_lazy env f
+    let x = compile_lazy (map incr env) x
     f # x # (Mkap ::)
-
   | Lam _ ->
       error "Can not compile lambda expression, did you forget to lift?"
+  | Case _ ->
+      error "Case expression in lazy context"
+  | Lit i -> (Push (Int i) ::)
+  | Let (vs, e) ->
+      compile_let compile_lazy env vs e
+
+and compile_strict (env : M.t string slot) = function
   | Case (sc, alts) ->
     let rec go_alts = function
       | [] -> []
@@ -172,27 +189,36 @@ let rec compile (env : M.t string slot) = function
             |> flip zip [Ls k | with k <- [c_arity - 1, c_arity - 2 .. 0]]
             |> M.from_list
             |> M.union (offs (c_arity + 1) <$> env)
-        (c_arity, compile env exp [Slide c_arity]) :: go_alts rest
-    compile env sc # (Eval ::) # (Casejump (go_alts alts) ::)
-  | Lit i -> (Push (Int i) ::)
-  | Let (vs, e) ->
-    let n = length vs
-    let env =
-      vs
-        |> map (fun (x, _) -> x)
-        |> flip zip [Ls x | with x <- [n - 1, n - 2 .. 0]]
-        |> M.from_list
-        |> M.union (offs n <$> env)
-    let defs = zip [1..n] vs
-    let rec go : list (int * string * expr) -> dlist gm_code = function
-      | [] -> id
-      | Cons ((k, (_, exp)), rest) ->
-          compile env exp # (Update (n - k) ::) # go rest
-    (Alloc n ::) # go defs # compile env e # (Slide n ::)
+        (c_arity, compile_strict env exp [Slide (c_arity + 1)]) :: go_alts rest
+    compile_strict env sc # (Casejump (go_alts alts) ::)
+  | 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_let cont env vs e =
+  let n = length vs
+  let env =
+    vs
+      |> map (fun (x, _) -> x)
+      |> flip zip [Ls x | with x <- [n - 1, n - 2 .. 0]]
+      |> M.from_list
+      |> M.union (offs n <$> env)
+  let defs = zip [1..n] vs
+  let rec go : list (int * string * expr) -> dlist gm_code = function
+    | [] -> id
+    | Cons ((k, (_, exp)), rest) ->
+        compile_lazy env exp # (Update (n - k) ::) # go rest
+  (Alloc n ::) # go defs # cont env e # (Slide n ::)
 
 let supercomb (_, args, exp) =
   let env = M.from_list (zip args [0..length args])
-  let k = compile (M.from_list (zip args (As <$> [0..length args]))) exp
+  let k = compile_strict (M.from_list (zip args (As <$> [0..length args]))) exp
   k [Update (length env), Pop (length env), Unwind]
 
 let compile_cons =
@@ -210,9 +236,17 @@ let compile_cons =
   go 0
 
 let program decs =
+  let rec globals s = function
+    | [] -> s
+    | Cons (Decl (n, _, _), r) -> globals (S.insert n s) r
+    | Cons (Data (_, _, c), r) ->
+        globals (foldl (fun s (Constr (n, _)) -> S.insert n s) s c) r
+    | Cons (Foreign (Fimport {var=n}), r) ->
+        globals (S.insert n s) r
+
   let (decs, (_, lams, _)) =
-    run_state (traverse (lambda_lift_sc # eta_contract) decs)
-      (0, [], S.empty)
+    run_state (traverse (map eta_contract # lambda_lift_sc) decs)
+      (0, [], globals S.empty decs)
 
   let define nm k =
     let! x = get
@@ -232,7 +266,9 @@ let program decs =
       | Data (_, _, cs) -> pure (compile_cons cs)
       | Foreign (Fimport { cc = Prim, var } as fi) ->
         define var (
-          let (code, _) = cg_prim fi
+          let (code, h) = cg_prim fi
+          print h
+          prim_scs := M.insert var h !prim_scs
           pure [code]
         )
       | Foreign f -> pure [Fd f]

driver.ml

@@ -19,6 +19,7 @@ let go infile outfile =
     match lex prog str with
     | Right (ds, _) ->
       ds
+        |> T.tc_program
         |> C.program
         |> A.assm_program
         |> write_bytes outfile

tc.ml

@@ -413,12 +413,14 @@ let rec add_missing_vars scope = function
 let tc_program (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, defs)
+    | [] -> (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?"
@@ -429,14 +431,16 @@ 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 ]
-        (dt_info, M.union (map force scope') val_scope, ty_scope, decs ++ defs)
+        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, reverse ds ++ out)
+          | [] -> (dt, vl, ty, ds)
           | Cons ((name, kind, constrs, info : list dt_info), rest) ->
               go
                 (foldl (fun i {name} -> M.insert name info i) dt info)
@@ -447,6 +451,6 @@ let tc_program (prog : list decl) =
                   vl info)
                 (Data (name, [], fix_constr <$> constrs) :: ds)
                 rest
-        go dt_info ty_scope val_scope [] r
+        go dt_info ty_scope val_scope out r
   let (_, _, _, p) = foldl tc_one (M.empty, M.empty, M.empty, []) plan
   p