Skip to content

Commit afce3aa

Browse files
ospencerospencer
authored
feat(compiler): Allow function re-exports to use regular call instruction (#1176)
* feat(compiler): Allow function re-exports to use regular call instruction * snapshots * PR feedback
1 parent b2d7440 commit afce3aa

File tree

130 files changed

+1967
-2954
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

130 files changed

+1967
-2954
lines changed

‎compiler/src/codegen/compcore.re‎

Lines changed: 32 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -3393,53 +3393,51 @@ let compile_imports = (wasm_mod, env, {imports}) => {
33933393
);
33943394
};
33953395

3396-
let compile_exports = (wasm_mod, env, {functions, imports, exports, globals}) => {
3397-
let compile_export = (i, {ex_name, ex_global_index}) => {
3398-
let internal_name = Printf.sprintf("global_%ld", ex_global_index);
3399-
let exported_name = "GRAIN$EXPORT$" ++ Ident.name(ex_name);
3400-
ignore @@ Export.add_global_export(wasm_mod, internal_name, exported_name);
3401-
};
3402-
3403-
let compile_external_function_export = ((internal_name, external_name)) => {
3404-
ignore @@
3405-
Export.add_function_export(wasm_mod, internal_name, external_name);
3396+
let compile_exports = (wasm_mod, env, {imports, exports, globals}) => {
3397+
let compile_export = (i, export) => {
3398+
switch (export) {
3399+
| GlobalExport({ex_global_name, ex_global_index}) =>
3400+
let internal_name = Printf.sprintf("global_%ld", ex_global_index);
3401+
let exported_name = "GRAIN$EXPORT$" ++ Ident.name(ex_global_name);
3402+
ignore @@
3403+
Export.add_global_export(wasm_mod, internal_name, exported_name);
3404+
| FunctionExport({ex_function_internal_name, ex_function_name}) =>
3405+
ignore @@
3406+
Export.add_function_export(
3407+
wasm_mod,
3408+
ex_function_internal_name,
3409+
ex_function_name,
3410+
)
3411+
};
34063412
};
34073413

34083414
let exports = {
3409-
let exported = Hashtbl.create(14);
3415+
module StringSet = Set.Make(String);
3416+
let exported_globals = ref(StringSet.empty);
3417+
let exported_functions = ref(StringSet.empty);
34103418
/* Exports are already reversed, so keeping the first of any name is the correct behavior. */
34113419
List.filter(
3412-
({ex_name}) =>
3413-
if (Hashtbl.mem(exported, Ident.name(ex_name))) {
3420+
fun
3421+
| GlobalExport({ex_global_name}) =>
3422+
if (StringSet.mem(Ident.name(ex_global_name), exported_globals^)) {
34143423
false;
34153424
} else {
3416-
Hashtbl.add(exported, Ident.name(ex_name), ());
3425+
exported_globals :=
3426+
StringSet.add(Ident.name(ex_global_name), exported_globals^);
3427+
true;
3428+
}
3429+
| FunctionExport({ex_function_name}) =>
3430+
if (StringSet.mem(ex_function_name, exported_functions^)) {
3431+
false;
3432+
} else {
3433+
exported_functions :=
3434+
StringSet.add(ex_function_name, exported_functions^);
34173435
true;
34183436
},
34193437
exports,
34203438
);
34213439
};
3422-
let functions = {
3423-
let exported = Hashtbl.create(14);
3424-
/* Functions will be reversed, so keeping the first of any name is the correct behavior. */
3425-
List.filter_map(
3426-
({index, name, id}) =>
3427-
switch (name) {
3428-
| Some(name) =>
3429-
if (Hashtbl.mem(exported, name)) {
3430-
None;
3431-
} else {
3432-
Hashtbl.add(exported, name, ());
3433-
let internal_name = Ident.unique_name(id);
3434-
Some((internal_name, name));
3435-
}
3436-
| None => None
3437-
},
3438-
functions,
3439-
);
3440-
};
34413440
List.iteri(compile_export, exports);
3442-
List.iter(compile_external_function_export, List.rev(functions));
34433441
ignore @@ Export.add_function_export(wasm_mod, grain_main, grain_main);
34443442
ignore @@ Export.add_function_export(wasm_mod, grain_start, grain_start);
34453443
ignore @@

‎compiler/src/codegen/mashtree.re‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -456,10 +456,15 @@ type import = {
456456
};
457457

458458
[@deriving sexp]
459-
type export = {
460-
ex_name: Ident.t,
461-
ex_global_index: int32,
462-
};
459+
type export =
460+
| FunctionExport({
461+
ex_function_name: string,
462+
ex_function_internal_name: string,
463+
})
464+
| GlobalExport({
465+
ex_global_name: Ident.t,
466+
ex_global_index: int32,
467+
});
463468

464469
[@deriving sexp]
465470
type mash_function = {

‎compiler/src/codegen/transl_anf.re‎

Lines changed: 63 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,9 @@ let global_index = ref(0);
7979
let global_exports = () => {
8080
let tbl = global_table^;
8181
Ident.fold_all(
82-
(ex_name, (exported, ex_global_index, _), acc) =>
82+
(ex_global_name, (exported, ex_global_index, _), acc) =>
8383
if (exported) {
84-
[{ex_name, ex_global_index}, ...acc];
84+
[GlobalExport({ex_global_name, ex_global_index}), ...acc];
8585
} else {
8686
acc;
8787
},
@@ -511,6 +511,12 @@ let run_register_allocation = (instrs: list(Mashtree.instr)) => {
511511
|> run(Types.StackAllocated(WasmF64));
512512
};
513513

514+
let grain_import_name = (mod_, name) =>
515+
Printf.sprintf("gimport_%s_%s", mod_, name);
516+
517+
let wasm_import_name = (mod_, name) =>
518+
Printf.sprintf("wimport_%s_%s", mod_, name);
519+
514520
let compile_const = (c: Asttypes.constant) =>
515521
switch (c) {
516522
| Const_number(Const_number_int(i)) => MConstI32(Int64.to_int32(i))
@@ -565,6 +571,7 @@ let known_function = f =>
565571
let compile_lambda =
566572
(~name=?, id, env, args, body, attrs, loc): Mashtree.closure_data => {
567573
register_function(Internal(id));
574+
568575
let (body, return_type) = body;
569576
// NOTE: we special-case `id`, since we want to
570577
// have simply-recursive uses of identifiers use
@@ -654,6 +661,7 @@ let compile_lambda =
654661
let compile_wrapper =
655662
(~export_name=?, id, env, func_name, args, rets): Mashtree.closure_data => {
656663
register_function(Internal(id));
664+
657665
let body = [
658666
{
659667
instr_desc:
@@ -1058,7 +1066,7 @@ let lift_imports = (env, imports) => {
10581066
| GrainValue(mod_, name) =>
10591067
let mimp_mod = Ident.create_persistent(mod_);
10601068
let mimp_name = Ident.create_persistent(name);
1061-
let import_name = Printf.sprintf("gimport_%s_%s", mod_, name);
1069+
let import_name = grain_import_name(mod_, name);
10621070
let (alloc, mods) =
10631071
switch (imp_shape) {
10641072
| GlobalShape(alloc) => (
@@ -1137,8 +1145,7 @@ let lift_imports = (env, imports) => {
11371145
Ident.add(
11381146
imp_use_id,
11391147
MGlobalBind(
1140-
Printf.sprintf(
1141-
"wimport_%s_%s",
1148+
wasm_import_name(
11421149
Ident.name(mimp_mod),
11431150
Ident.name(mimp_name),
11441151
),
@@ -1160,7 +1167,7 @@ let lift_imports = (env, imports) => {
11601167
mimp_setup: MWrap(Int32.zero),
11611168
mimp_used: true,
11621169
};
1163-
let func_name = Printf.sprintf("wimport_%s_%s", mod_, name);
1170+
let func_name = wasm_import_name(mod_, name);
11641171
let export_name =
11651172
if (exported) {
11661173
Some(name);
@@ -1228,44 +1235,78 @@ let lift_imports = (env, imports) => {
12281235
(imports, setups, env);
12291236
};
12301237

1231-
let transl_signature = (functions, signature) => {
1238+
let transl_signature = (~functions, ~imports, signature) => {
12321239
open Types;
12331240

1241+
let function_exports = ref([]);
1242+
12341243
// At this point in compilation, we know which functions can be called
12351244
// directly/indirectly at the wasm level. We add this information to the
12361245
// module signature.
12371246
let func_map = Ident_tbl.create(30);
12381247
List.iter(
1239-
func => Ident_tbl.add(func_map, (func: mash_function).id, func),
1248+
(func: mash_function) =>
1249+
switch (func.name) {
1250+
| Some(name) =>
1251+
Ident_tbl.add(func_map, func.id, Ident.unique_name(func.id))
1252+
| None => ()
1253+
},
12401254
functions,
12411255
);
1256+
List.iter(
1257+
imp =>
1258+
switch (imp.imp_shape) {
1259+
| FunctionShape(_) =>
1260+
let internal_name =
1261+
switch (imp.imp_desc) {
1262+
| GrainValue(mod_, name) => grain_import_name(mod_, name)
1263+
| WasmFunction(mod_, name) => Ident.unique_name(imp.imp_use_id)
1264+
| _ => failwith("Impossible: Wasm or js value had FunctionShape")
1265+
};
1266+
Ident_tbl.add(func_map, imp.imp_use_id, internal_name);
1267+
| _ => ()
1268+
},
1269+
imports,
1270+
);
12421271
let sign =
12431272
List.map(
12441273
fun
12451274
| TSigValue(
1246-
_,
1275+
vid,
12471276
{
12481277
val_repr: ReprFunction(args, rets, _),
12491278
val_fullpath: Path.PIdent(id),
12501279
} as vd,
12511280
) => {
12521281
switch (Ident_tbl.find_opt(func_map, id)) {
1253-
| Some({name: Some(name)}) =>
1282+
| Some(internal_name) =>
1283+
let external_name = Ident.name(vid);
1284+
function_exports :=
1285+
[
1286+
FunctionExport({
1287+
ex_function_name: external_name,
1288+
ex_function_internal_name: internal_name,
1289+
}),
1290+
...function_exports^,
1291+
];
12541292
TSigValue(
1255-
id,
1256-
{...vd, val_repr: ReprFunction(args, rets, Direct(name))},
1257-
)
1293+
vid,
1294+
{
1295+
...vd,
1296+
val_repr: ReprFunction(args, rets, Direct(external_name)),
1297+
},
1298+
);
12581299
| _ =>
12591300
TSigValue(
1260-
id,
1301+
vid,
12611302
{...vd, val_repr: ReprFunction(args, rets, Indirect)},
12621303
)
12631304
};
12641305
}
12651306
| _ as item => item,
12661307
signature.Cmi_format.cmi_sign,
12671308
);
1268-
{...signature, cmi_sign: sign};
1309+
({...signature, cmi_sign: sign}, function_exports^);
12691310
};
12701311

12711312
let transl_anf_program =
@@ -1297,14 +1338,19 @@ let transl_anf_program =
12971338
};
12981339
let main_body =
12991340
run_register_allocation @@ setups @ compile_anf_expr(env, anf_prog.body);
1300-
let exports = global_exports();
13011341
let functions =
13021342
List.map(
13031343
({body} as f: Mashtree.mash_function) =>
13041344
{...f, body: run_register_allocation(body)},
13051345
compile_remaining_worklist(),
13061346
);
1307-
let signature = transl_signature(functions, anf_prog.signature);
1347+
let (signature, function_exports) =
1348+
transl_signature(
1349+
~functions,
1350+
~imports=anf_prog.imports,
1351+
anf_prog.signature,
1352+
);
1353+
let exports = function_exports @ global_exports();
13081354

13091355
{
13101356
functions,

‎compiler/src/middle_end/optimize_constants.re‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,30 @@ module ConstantPropagationArg: Anf_mapper.MapArgument = {
4949
}
5050
| _ => i
5151
};
52+
53+
let leave_anf_program = ({signature} as p) => {
54+
// Support directly exporting imports
55+
let cmi_sign =
56+
List.map(
57+
fun
58+
| TSigValue(vid, {val_fullpath: Path.PIdent(id)} as vd) as item => {
59+
switch (Ident.find_same_opt(id, known_constants^)) {
60+
| Some(ImmId(immid)) =>
61+
TSigValue(vid, {...vd, val_fullpath: Path.PIdent(immid)})
62+
| _ => item
63+
};
64+
}
65+
| item => item,
66+
signature.cmi_sign,
67+
);
68+
{
69+
...p,
70+
signature: {
71+
...signature,
72+
cmi_sign,
73+
},
74+
};
75+
};
5276
};
5377

5478
module ConstantPropagationMapper = Anf_mapper.MakeMap(ConstantPropagationArg);

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

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,28 @@
11
basic functionality › modulo4
22
(module
33
(type $none_=>_i32 (func (result i32)))
4-
(type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
54
(type $none_=>_none (func))
65
(type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
6+
(type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
77
(import \"_grainEnv\" \"mem\" (memory $0 0))
8-
(import \"_grainEnv\" \"tbl\" (table $tbl 0 funcref))
98
(import \"_grainEnv\" \"relocBase\" (global $wimport__grainEnv_relocBase i32))
109
(import \"GRAIN$MODULE$runtime/gc\" \"GRAIN$EXPORT$incRef\" (global $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef (mut i32)))
1110
(import \"GRAIN$MODULE$pervasives\" \"GRAIN$EXPORT$%\" (global $gimport_pervasives_% (mut i32)))
1211
(import \"GRAIN$MODULE$runtime/gc\" \"incRef\" (func $wimport_GRAIN$MODULE$runtime/gc_incRef (param i32 i32) (result i32)))
12+
(import \"GRAIN$MODULE$pervasives\" \"%\" (func $gimport_pervasives_% (param i32 i32 i32) (result i32)))
1313
(global $global_1 i32 (i32.const 0))
1414
(export \"memory\" (memory $0))
1515
(export \"_gmain\" (func $_gmain))
1616
(export \"_start\" (func $_start))
1717
(export \"GRAIN$TABLE_SIZE\" (global $global_1))
1818
(func $_gmain (; has Stack IR ;) (result i32)
19-
(local $0 i32)
20-
(call_indirect (type $i32_i32_i32_=>_i32)
21-
(local.tee $0
22-
(call $wimport_GRAIN$MODULE$runtime/gc_incRef
23-
(global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
24-
(global.get $gimport_pervasives_%)
25-
)
19+
(call $gimport_pervasives_%
20+
(call $wimport_GRAIN$MODULE$runtime/gc_incRef
21+
(global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
22+
(global.get $gimport_pervasives_%)
2623
)
2724
(i32.const -33)
2825
(i32.const 35)
29-
(i32.load offset=8
30-
(local.get $0)
31-
)
3226
)
3327
)
3428
(func $_start (; has Stack IR ;)

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

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,28 @@
11
basic functionality › land4
22
(module
33
(type $none_=>_i32 (func (result i32)))
4-
(type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
54
(type $none_=>_none (func))
65
(type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
6+
(type $i32_i32_i32_=>_i32 (func (param i32 i32 i32) (result i32)))
77
(import \"_grainEnv\" \"mem\" (memory $0 0))
8-
(import \"_grainEnv\" \"tbl\" (table $tbl 0 funcref))
98
(import \"_grainEnv\" \"relocBase\" (global $wimport__grainEnv_relocBase i32))
109
(import \"GRAIN$MODULE$runtime/gc\" \"GRAIN$EXPORT$incRef\" (global $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef (mut i32)))
1110
(import \"GRAIN$MODULE$pervasives\" \"GRAIN$EXPORT$&\" (global $gimport_pervasives_& (mut i32)))
1211
(import \"GRAIN$MODULE$runtime/gc\" \"incRef\" (func $wimport_GRAIN$MODULE$runtime/gc_incRef (param i32 i32) (result i32)))
12+
(import \"GRAIN$MODULE$pervasives\" \"&\" (func $gimport_pervasives_& (param i32 i32 i32) (result i32)))
1313
(global $global_1 i32 (i32.const 0))
1414
(export \"memory\" (memory $0))
1515
(export \"_gmain\" (func $_gmain))
1616
(export \"_start\" (func $_start))
1717
(export \"GRAIN$TABLE_SIZE\" (global $global_1))
1818
(func $_gmain (; has Stack IR ;) (result i32)
19-
(local $0 i32)
20-
(call_indirect (type $i32_i32_i32_=>_i32)
21-
(local.tee $0
22-
(call $wimport_GRAIN$MODULE$runtime/gc_incRef
23-
(global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
24-
(global.get $gimport_pervasives_&)
25-
)
19+
(call $gimport_pervasives_&
20+
(call $wimport_GRAIN$MODULE$runtime/gc_incRef
21+
(global.get $wimport_GRAIN$MODULE$runtime/gc_GRAIN$EXPORT$incRef)
22+
(global.get $gimport_pervasives_&)
2623
)
2724
(i32.const 1)
2825
(i32.const 1)
29-
(i32.load offset=8
30-
(local.get $0)
31-
)
3226
)
3327
)
3428
(func $_start (; has Stack IR ;)

0 commit comments

Comments
 (0)
⚡