feat(compiler): Allow function re-exports to use regular call instruc…

…tion (#1176)

* feat(compiler): Allow function re-exports to use regular call instruction

* snapshots

* PR feedback

compiler/src/codegen/compcore.re

@@ -3393,53 +3393,51 @@ let compile_imports = (wasm_mod, env, {imports}) => {
   );
 };
 
-let compile_exports = (wasm_mod, env, {functions, imports, exports, globals}) => {
-  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);
-    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 compile_exports = (wasm_mod, env, {imports, exports, globals}) => {
+  let compile_export = (i, export) => {
+    switch (export) {
+    | GlobalExport({ex_global_name, ex_global_index}) =>
+      let internal_name = Printf.sprintf("global_%ld", ex_global_index);
+      let exported_name = "GRAIN$EXPORT$" ++ Ident.name(ex_global_name);
+      ignore @@
+      Export.add_global_export(wasm_mod, internal_name, exported_name);
+    | FunctionExport({ex_function_internal_name, ex_function_name}) =>
+      ignore @@
+      Export.add_function_export(
+        wasm_mod,
+        ex_function_internal_name,
+        ex_function_name,
+      )
+    };
   };
 
   let exports = {
-    let exported = Hashtbl.create(14);
+    module StringSet = Set.Make(String);
+    let exported_globals = ref(StringSet.empty);
+    let exported_functions = ref(StringSet.empty);
     /* Exports are already reversed, so keeping the first of any name is the correct behavior. */
     List.filter(
-      ({ex_name}) =>
-        if (Hashtbl.mem(exported, Ident.name(ex_name))) {
+      fun
+      | GlobalExport({ex_global_name}) =>
+        if (StringSet.mem(Ident.name(ex_global_name), exported_globals^)) {
           false;
         } else {
-          Hashtbl.add(exported, Ident.name(ex_name), ());
+          exported_globals :=
+            StringSet.add(Ident.name(ex_global_name), exported_globals^);
+          true;
+        }
+      | FunctionExport({ex_function_name}) =>
+        if (StringSet.mem(ex_function_name, exported_functions^)) {
+          false;
+        } else {
+          exported_functions :=
+            StringSet.add(ex_function_name, exported_functions^);
           true;
         },
       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, id}) =>
-        switch (name) {
-        | Some(name) =>
-          if (Hashtbl.mem(exported, name)) {
-            None;
-          } else {
-            Hashtbl.add(exported, name, ());
-            let internal_name = Ident.unique_name(id);
-            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, grain_main, grain_main);
   ignore @@ Export.add_function_export(wasm_mod, grain_start, grain_start);
   ignore @@

compiler/src/codegen/mashtree.re

@@ -456,10 +456,15 @@ type import = {
 };
 
 [@deriving sexp]
-type export = {
-  ex_name: Ident.t,
-  ex_global_index: int32,
-};
+type export =
+  | FunctionExport({
+      ex_function_name: string,
+      ex_function_internal_name: string,
+    })
+  | GlobalExport({
+      ex_global_name: Ident.t,
+      ex_global_index: int32,
+    });
 
 [@deriving sexp]
 type mash_function = {

compiler/src/codegen/transl_anf.re

@@ -79,9 +79,9 @@ let global_index = ref(0);
 let global_exports = () => {
   let tbl = global_table^;
   Ident.fold_all(
-    (ex_name, (exported, ex_global_index, _), acc) =>
+    (ex_global_name, (exported, ex_global_index, _), acc) =>
       if (exported) {
-        [{ex_name, ex_global_index}, ...acc];
+        [GlobalExport({ex_global_name, ex_global_index}), ...acc];
       } else {
         acc;
       },
@@ -511,6 +511,12 @@ let run_register_allocation = (instrs: list(Mashtree.instr)) => {
   |> run(Types.StackAllocated(WasmF64));
 };
 
+let grain_import_name = (mod_, name) =>
+  Printf.sprintf("gimport_%s_%s", mod_, name);
+
+let wasm_import_name = (mod_, name) =>
+  Printf.sprintf("wimport_%s_%s", mod_, name);
+
 let compile_const = (c: Asttypes.constant) =>
   switch (c) {
   | Const_number(Const_number_int(i)) => MConstI32(Int64.to_int32(i))
@@ -565,6 +571,7 @@ let known_function = f =>
 let compile_lambda =
     (~name=?, id, env, args, body, attrs, loc): Mashtree.closure_data => {
   register_function(Internal(id));
+
   let (body, return_type) = body;
   // NOTE: we special-case `id`, since we want to
   //       have simply-recursive uses of identifiers use
@@ -654,6 +661,7 @@ let compile_lambda =
 let compile_wrapper =
     (~export_name=?, id, env, func_name, args, rets): Mashtree.closure_data => {
   register_function(Internal(id));
+
   let body = [
     {
       instr_desc:
@@ -1058,7 +1066,7 @@ let lift_imports = (env, imports) => {
     | GrainValue(mod_, name) =>
       let mimp_mod = Ident.create_persistent(mod_);
       let mimp_name = Ident.create_persistent(name);
-      let import_name = Printf.sprintf("gimport_%s_%s", mod_, name);
+      let import_name = grain_import_name(mod_, name);
       let (alloc, mods) =
         switch (imp_shape) {
         | GlobalShape(alloc) => (
@@ -1137,8 +1145,7 @@ let lift_imports = (env, imports) => {
             Ident.add(
               imp_use_id,
               MGlobalBind(
-                Printf.sprintf(
-                  "wimport_%s_%s",
+                wasm_import_name(
                   Ident.name(mimp_mod),
                   Ident.name(mimp_name),
                 ),
@@ -1160,7 +1167,7 @@ let lift_imports = (env, imports) => {
         mimp_setup: MWrap(Int32.zero),
         mimp_used: true,
       };
-      let func_name = Printf.sprintf("wimport_%s_%s", mod_, name);
+      let func_name = wasm_import_name(mod_, name);
       let export_name =
         if (exported) {
           Some(name);
@@ -1228,44 +1235,78 @@ let lift_imports = (env, imports) => {
   (imports, setups, env);
 };
 
-let transl_signature = (functions, signature) => {
+let transl_signature = (~functions, ~imports, signature) => {
   open Types;
 
+  let function_exports = ref([]);
+
   // At this point in compilation, we know which functions can be called
   // directly/indirectly at the wasm level. We add this information to the
   // module signature.
   let func_map = Ident_tbl.create(30);
   List.iter(
-    func => Ident_tbl.add(func_map, (func: mash_function).id, func),
+    (func: mash_function) =>
+      switch (func.name) {
+      | Some(name) =>
+        Ident_tbl.add(func_map, func.id, Ident.unique_name(func.id))
+      | None => ()
+      },
     functions,
   );
+  List.iter(
+    imp =>
+      switch (imp.imp_shape) {
+      | FunctionShape(_) =>
+        let internal_name =
+          switch (imp.imp_desc) {
+          | GrainValue(mod_, name) => grain_import_name(mod_, name)
+          | WasmFunction(mod_, name) => Ident.unique_name(imp.imp_use_id)
+          | _ => failwith("Impossible: Wasm or js value had FunctionShape")
+          };
+        Ident_tbl.add(func_map, imp.imp_use_id, internal_name);
+      | _ => ()
+      },
+    imports,
+  );
   let sign =
     List.map(
       fun
       | TSigValue(
-          _,
+          vid,
           {
             val_repr: ReprFunction(args, rets, _),
             val_fullpath: Path.PIdent(id),
           } as vd,
         ) => {
           switch (Ident_tbl.find_opt(func_map, id)) {
-          | Some({name: Some(name)}) =>
+          | Some(internal_name) =>
+            let external_name = Ident.name(vid);
+            function_exports :=
+              [
+                FunctionExport({
+                  ex_function_name: external_name,
+                  ex_function_internal_name: internal_name,
+                }),
+                ...function_exports^,
+              ];
             TSigValue(
-              id,
-              {...vd, val_repr: ReprFunction(args, rets, Direct(name))},
-            )
+              vid,
+              {
+                ...vd,
+                val_repr: ReprFunction(args, rets, Direct(external_name)),
+              },
+            );
           | _ =>
             TSigValue(
-              id,
+              vid,
               {...vd, val_repr: ReprFunction(args, rets, Indirect)},
             )
           };
         }
       | _ as item => item,
       signature.Cmi_format.cmi_sign,
     );
-  {...signature, cmi_sign: sign};
+  ({...signature, cmi_sign: sign}, function_exports^);
 };
 
 let transl_anf_program =
@@ -1297,14 +1338,19 @@ let transl_anf_program =
   };
   let main_body =
     run_register_allocation @@ setups @ compile_anf_expr(env, anf_prog.body);
-  let exports = global_exports();
   let functions =
     List.map(
       ({body} as f: Mashtree.mash_function) =>
         {...f, body: run_register_allocation(body)},
       compile_remaining_worklist(),
     );
-  let signature = transl_signature(functions, anf_prog.signature);
+  let (signature, function_exports) =
+    transl_signature(
+      ~functions,
+      ~imports=anf_prog.imports,
+      anf_prog.signature,
+    );
+  let exports = function_exports @ global_exports();
 
   {
     functions,

compiler/src/middle_end/optimize_constants.re

@@ -49,6 +49,30 @@ module ConstantPropagationArg: Anf_mapper.MapArgument = {
       }
     | _ => i
     };
+
+  let leave_anf_program = ({signature} as p) => {
+    // Support directly exporting imports
+    let cmi_sign =
+      List.map(
+        fun
+        | TSigValue(vid, {val_fullpath: Path.PIdent(id)} as vd) as item => {
+            switch (Ident.find_same_opt(id, known_constants^)) {
+            | Some(ImmId(immid)) =>
+              TSigValue(vid, {...vd, val_fullpath: Path.PIdent(immid)})
+            | _ => item
+            };
+          }
+        | item => item,
+        signature.cmi_sign,
+      );
+    {
+      ...p,
+      signature: {
+        ...signature,
+        cmi_sign,
+      },
+    };
+  };
 };
 
 module ConstantPropagationMapper = Anf_mapper.MakeMap(ConstantPropagationArg);

compiler/test/__snapshots__/basic_functionality.0996c5f7.0.snapshot

@@ -1,34 +1,28 @@
 basic functionality › modulo4
 (module
  (type $none_=>_i32 (func (result i32)))
- (type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
  (type $none_=>_none (func))
  (type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
+ (type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
  (import \"_grainEnv\" \"mem\" (memory $0 0))
- (import \"_grainEnv\" \"tbl\" (table $tbl 0 funcref))
  (import \"_grainEnv\" \"relocBase\" (global $wimport__grainEnv_relocBase i32))
  (import \"GRAIN$MODULE$runtime/gc\" \"GRAIN$EXPORT$incRef\" (global $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef (mut i32)))
  (import \"GRAIN$MODULE$pervasives\" \"GRAIN$EXPORT$%\" (global $gimport_pervasives_% (mut i32)))
  (import \"GRAIN$MODULE$runtime/gc\" \"incRef\" (func $wimport_GRAIN$MODULE$runtime/gc_incRef (param i32 i32) (result i32)))
+ (import \"GRAIN$MODULE$pervasives\" \"%\" (func $gimport_pervasives_% (param i32 i32 i32) (result i32)))
  (global $global_1 i32 (i32.const 0))
  (export \"memory\" (memory $0))
  (export \"_gmain\" (func $_gmain))
  (export \"_start\" (func $_start))
  (export \"GRAIN$TABLE_SIZE\" (global $global_1))
  (func $_gmain (; has Stack IR ;) (result i32)
-  (local $0 i32)
-  (call_indirect (type $i32_i32_i32_=>_i32)
-   (local.tee $0
-    (call $wimport_GRAIN$MODULE$runtime/gc_incRef
-     (global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
-     (global.get $gimport_pervasives_%)
-    )
+  (call $gimport_pervasives_%
+   (call $wimport_GRAIN$MODULE$runtime/gc_incRef
+    (global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
+    (global.get $gimport_pervasives_%)
    )
    (i32.const -33)
    (i32.const 35)
-   (i32.load offset=8
-    (local.get $0)
-   )
   )
  )
  (func $_start (; has Stack IR ;)

compiler/test/__snapshots__/basic_functionality.0a230f18.0.snapshot

@@ -1,34 +1,28 @@
 basic functionality › land4
 (module
  (type $none_=>_i32 (func (result i32)))
- (type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
  (type $none_=>_none (func))
  (type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
+ (type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
  (import \"_grainEnv\" \"mem\" (memory $0 0))
- (import \"_grainEnv\" \"tbl\" (table $tbl 0 funcref))
  (import \"_grainEnv\" \"relocBase\" (global $wimport__grainEnv_relocBase i32))
  (import \"GRAIN$MODULE$runtime/gc\" \"GRAIN$EXPORT$incRef\" (global $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef (mut i32)))
  (import \"GRAIN$MODULE$pervasives\" \"GRAIN$EXPORT$&\" (global $gimport_pervasives_& (mut i32)))
  (import \"GRAIN$MODULE$runtime/gc\" \"incRef\" (func $wimport_GRAIN$MODULE$runtime/gc_incRef (param i32 i32) (result i32)))
+ (import \"GRAIN$MODULE$pervasives\" \"&\" (func $gimport_pervasives_& (param i32 i32 i32) (result i32)))
  (global $global_1 i32 (i32.const 0))
  (export \"memory\" (memory $0))
  (export \"_gmain\" (func $_gmain))
  (export \"_start\" (func $_start))
  (export \"GRAIN$TABLE_SIZE\" (global $global_1))
  (func $_gmain (; has Stack IR ;) (result i32)
-  (local $0 i32)
-  (call_indirect (type $i32_i32_i32_=>_i32)
-   (local.tee $0
-    (call $wimport_GRAIN$MODULE$runtime/gc_incRef
-     (global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
-     (global.get $gimport_pervasives_&)
-    )
+  (call $gimport_pervasives_&
+   (call $wimport_GRAIN$MODULE$runtime/gc_incRef
+    (global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
+    (global.get $gimport_pervasives_&)
    )
    (i32.const 1)
    (i32.const 1)
-   (i32.load offset=8
-    (local.get $0)
-   )
   )
  )
  (func $_start (; has Stack IR ;)