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

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 ;)