feat: Export compiled functions from modules (#529)

compiler/src/codegen/compcore.re

@@ -3395,12 +3395,15 @@ let compile_imports = (wasm_mod, env, {imports}) => {
 };
 
 let compile_exports = (wasm_mod, env, {functions, imports, exports, globals}) => {
-  let compile_exports = (i, {ex_name, ex_global_index}) => {
+  let compile_export = (i, {ex_name, ex_global_index}) => {
     let internal_name = Printf.sprintf("global_%ld", ex_global_index);
     let exported_name = "GRAIN$EXPORT$" ++ Ident.name(ex_name);
-    let export =
-      Export.add_global_export(wasm_mod, internal_name, exported_name);
-    export;
+    ignore @@ Export.add_global_export(wasm_mod, internal_name, exported_name);
+  };
+
+  let compile_external_function_export = ((internal_name, external_name)) => {
+    ignore @@
+    Export.add_function_export(wasm_mod, internal_name, external_name);
   };
 
   let exports = {
@@ -3417,7 +3420,27 @@ let compile_exports = (wasm_mod, env, {functions, imports, exports, globals}) =>
       exports,
     );
   };
-  ignore @@ List.mapi(compile_exports, exports);
+  let functions = {
+    let exported = Hashtbl.create(14);
+    /* Functions will be reversed, so keeping the first of any name is the correct behavior. */
+    List.filter_map(
+      ({index, name}) =>
+        switch (name) {
+        | Some(name) =>
+          if (Hashtbl.mem(exported, name)) {
+            None;
+          } else {
+            Hashtbl.add(exported, name, ());
+            let internal_name = Printf.sprintf("func_%ld", index);
+            Some((internal_name, name));
+          }
+        | None => None
+        },
+      functions,
+    );
+  };
+  List.iteri(compile_export, exports);
+  List.iter(compile_external_function_export, List.rev(functions));
   ignore @@ Export.add_function_export(wasm_mod, "_start", "_start");
   ignore @@
   Export.add_global_export(
@@ -3481,6 +3504,7 @@ let compile_main = (wasm_mod, env, prog) => {
     env,
     {
       index: Int32.of_int(-99),
+      name: Some("_start"),
       args: [],
       return_type: I32Type,
       body: prog.main_body,

compiler/src/codegen/mashtree.re

@@ -460,6 +460,7 @@ type export = {
 [@deriving sexp]
 type mash_function = {
   index: int32,
+  name: option(string),
   args: list(asmtype),
   return_type: asmtype,
   body: block,

compiler/src/codegen/transl_anf.re

@@ -38,6 +38,7 @@ type worklist_elt = {
   args: list(Types.allocation_type),
   return_type: Types.allocation_type,
   idx: int, /* Lambda-lifted index */
+  name: option(string),
   attrs: attributes,
   stack_size: Mashtree.stack_size,
   loc: Location.t,
@@ -477,7 +478,8 @@ let compile_imm = (env, i: imm_expression) =>
   | ImmId(id) => MImmBinding(find_id(id, env))
   };
 
-let compile_lambda = (env, args, body, attrs, loc): Mashtree.closure_data => {
+let compile_lambda =
+    (~name=?, env, args, body, attrs, loc): Mashtree.closure_data => {
   let (body, return_type) = body;
   let used_var_set = Anf_utils.anf_free_vars(body);
   let free_var_set =
@@ -525,6 +527,7 @@ let compile_lambda = (env, args, body, attrs, loc): Mashtree.closure_data => {
     body: Anf(body),
     env: lam_env,
     idx,
+    name,
     args,
     return_type,
     attrs,
@@ -582,6 +585,7 @@ let compile_wrapper = (env, func_name, args, ret): Mashtree.closure_data => {
     body: Precompiled(body),
     env: lam_env,
     idx,
+    name: None,
     args: [Types.HeapAllocated, ...args],
     return_type: Types.HeapAllocated,
     stack_size: {
@@ -732,10 +736,11 @@ let rec compile_comp = (env, c) => {
         MRecordSet(idx, compile_imm(env, arg)),
         compile_imm(env, record),
       )
-    | CLambda(args, body) =>
+    | CLambda(name, args, body) =>
       MAllocate(
         MClosure(
           compile_lambda(
+            ~name?,
             env,
             args,
             body,
@@ -893,11 +898,12 @@ let compile_remaining_worklist = () => {
   let compile_one =
       (
         funcs,
-        {idx: index, args, return_type, stack_size, attrs, loc} as cur: worklist_elt,
+        {idx: index, name, args, return_type, stack_size, attrs, loc} as cur: worklist_elt,
       ) => {
     let body = compile_worklist_elt(cur);
     let func = {
       index: Int32.of_int(index),
+      name,
       args: List.map(asmtype_of_alloctype, args),
       return_type: asmtype_of_alloctype(return_type),
       body,

compiler/src/middle_end/analyze_purity.re

@@ -139,7 +139,7 @@ let rec analyze_comp_expression =
       List.for_all(x => x, branches_purities);
     | CApp(_) => false
     | CAppBuiltin(_) => false
-    | CLambda(_, (body, _)) =>
+    | CLambda(_, _, (body, _)) =>
       analyze_anf_expression(body);
       true;
     | CNumber(_)

compiler/src/middle_end/analyze_tail_calls.re

@@ -39,7 +39,7 @@ let rec analyze_comp_expression =
     /* While this loop itself is not in tail position, we still want to analyze the body. */
     ignore @@ analyze_anf_expression(false, body);
     false;
-  | CLambda(args, (body, _)) =>
+  | CLambda(_, args, (body, _)) =>
     /* While this lambda itself is not in tail position, we still want to analyze the body. */
     ignore @@ analyze_anf_expression(true, body);
     false;
@@ -100,7 +100,7 @@ and analyze_anf_expression =
     List.iter(
       ((_, {comp_desc, comp_analyses} as bind)) =>
         switch (comp_desc) {
-        | CLambda(args, (body, _)) =>
+        | CLambda(_, args, (body, _)) =>
           if (analyze_anf_expression(true, body)) {
             push_tail_recursive(comp_analyses);
           }

compiler/src/middle_end/anf_helper.re

@@ -262,13 +262,13 @@ module Comp = {
       ~env?,
       CAppBuiltin(modname, name, args),
     );
-  let lambda = (~loc=?, ~attributes=?, ~env=?, args, body) =>
+  let lambda = (~loc=?, ~attributes=?, ~env=?, ~name=?, args, body) =>
     mk(
       ~loc?,
       ~attributes?,
       ~allocation_type=HeapAllocated,
       ~env?,
-      CLambda(args, body),
+      CLambda(name, args, body),
     );
   let string = (~loc=?, ~attributes=?, ~env=?, s) =>
     mk(

compiler/src/middle_end/anf_helper.rei

@@ -294,6 +294,7 @@ module Comp: {
       ~loc: loc=?,
       ~attributes: attributes=?,
       ~env: env=?,
+      ~name: string=?,
       list((ident, allocation_type)),
       (anf_expression, allocation_type)
     ) =>

compiler/src/middle_end/anf_iterator.re

@@ -97,7 +97,7 @@ module MakeIter = (Iter: IterArgument) => {
       iter_imm_expression(f);
       List.iter(iter_imm_expression, args);
     | CAppBuiltin(_, _, args) => List.iter(iter_imm_expression, args)
-    | CLambda(idents, (expr, _)) => iter_anf_expression(expr)
+    | CLambda(_, idents, (expr, _)) => iter_anf_expression(expr)
     | CString(s) => ()
     | CChar(c) => ()
     | CNumber(i) => ()

compiler/src/middle_end/anf_mapper.re

@@ -124,9 +124,9 @@ module MakeMap = (Iter: MapArgument) => {
         CApp((f, fty), args, tail);
       | CAppBuiltin(mod_, f, args) =>
         CAppBuiltin(mod_, f, List.map(map_imm_expression, args))
-      | CLambda(idents, (expr, alloc_ty)) =>
+      | CLambda(name, idents, (expr, alloc_ty)) =>
         let expr = map_anf_expression(expr);
-        CLambda(idents, (expr, alloc_ty));
+        CLambda(name, idents, (expr, alloc_ty));
       | CString(s) => CString(s)
       | CChar(c) => CChar(c)
       | CNumber(i) => CNumber(i)

compiler/src/middle_end/anf_utils.re

@@ -35,7 +35,7 @@ let rec anf_free_vars_help = (env, a: anf_expression) =>
 
 and comp_free_vars_help = (env, c: comp_expression) =>
   switch (c.comp_desc) {
-  | CLambda(args, (body, _)) =>
+  | CLambda(_, args, (body, _)) =>
     anf_free_vars_help(
       Ident.Set.union(
         env,

compiler/src/middle_end/anftree.re

@@ -340,6 +340,7 @@ and comp_expression_desc =
     )
   | CAppBuiltin(string, string, list(imm_expression))
   | CLambda(
+      option(string),
       list((Ident.t, allocation_type)),
       (anf_expression, allocation_type),
     )

compiler/src/middle_end/anftree.rei

@@ -325,6 +325,7 @@ and comp_expression_desc =
     )
   | CAppBuiltin(string, string, list(imm_expression))
   | CLambda(
+      option(string),
       list((Ident.t, allocation_type)),
       (anf_expression, allocation_type),
     )

compiler/src/middle_end/linearize.re

@@ -764,6 +764,7 @@ and bind_patts =
 
 and transl_comp_expression =
     (
+      ~name=?,
       {
         exp_desc,
         exp_type,
@@ -962,6 +963,7 @@ and transl_comp_expression =
         ~loc,
         ~attributes,
         ~env,
+        ~name?,
         [],
         (anf_body, get_allocation_type(body.exp_env, body.exp_type)),
       ),
@@ -1007,6 +1009,7 @@ and transl_comp_expression =
           ~loc,
           ~attributes,
           ~env,
+          ~name?,
           anf_args,
           (anf_body, get_allocation_type(body.exp_env, body.exp_type)),
         ),
@@ -1175,14 +1178,23 @@ let rec transl_anf_statement =
       ...vb_expr,
       exp_attributes: attributes @ vb_expr.exp_attributes,
     };
-    let (exp_ans, exp_setup) = transl_comp_expression(vb_expr);
+    let exported = export_flag == Exported;
+    let name =
+      if (exported) {
+        switch (vb_pat.pat_desc) {
+        | TPatVar(bind, _) => Some(Ident.name(bind))
+        | _ => None
+        };
+      } else {
+        None;
+      };
+    let (exp_ans, exp_setup) = transl_comp_expression(~name?, vb_expr);
     let (rest_setup, rest_imp) =
       transl_anf_statement({
         ...s,
         ttop_desc: TTopLet(export_flag, Nonrecursive, mut_flag, rest),
       });
     let rest_setup = Option.value(~default=[], rest_setup);
-    let exported = export_flag == Exported;
     let setup =
       switch (vb_pat.pat_desc) {
       | TPatVar(bind, _)