feat(compiler): Use data segments for string data

Author: ospencer

compiler/src/codegen/compcore.re

@@ -34,6 +34,17 @@ let gensym_label = s => {
 };
 let reset_labels = () => gensym_counter := 0;
 
+let data_segments: ref((int, list(Memory.segment))) = ref((0, []));
+let push_data_segment = segment => {
+  let (count, segments) = data_segments^;
+  data_segments := (count + 1, [segment, ...segments]);
+};
+let get_data_segment_count = () => {
+  let (count, _) = data_segments^;
+  count;
+};
+let reset_data_segments = () => data_segments := (0, []);
+
 /* Number of swap variables to allocate */
 let swap_slots_i32 = [|Type.int32, Type.int32, Type.int32|];
 let swap_slots_i64 = [|Type.int64|];
@@ -117,6 +128,7 @@ let rec resolve_func = (~env, name) => {
 
 let reset = () => {
   reset_labels();
+  reset_data_segments();
 };
 
 let get_runtime_heap_start = wasm_mod =>
@@ -1267,45 +1279,93 @@ let call_lambda =
 };
 
 let allocate_byte_like_from_buffer = (wasm_mod, env, buf, tag, label) => {
-  let ints_to_push: list(int64) = buf_to_ints(buf);
   let get_swap = () => get_swap(wasm_mod, env, 0);
   let tee_swap = tee_swap(wasm_mod, env, 0);
-  let preamble = [
-    store(
-      ~offset=0,
-      wasm_mod,
-      tee_swap(
-        heap_allocate(wasm_mod, env, 2 + 2 * List.length(ints_to_push)),
-      ),
-      Expression.Const.make(
-        wasm_mod,
-        const_int32(tag_val_of_heap_tag_type(tag)),
-      ),
-    ),
-    store(
-      ~offset=4,
+  if (Config.bulk_memory^) {
+    let segment_offset = get_data_segment_count();
+    let data_size = Buffer.length(buf);
+    Memory.(
+      push_data_segment({
+        data: Buffer.to_bytes(buf),
+        kind: Passive,
+        size: data_size,
+      })
+    );
+    Expression.Block.make(
       wasm_mod,
-      get_swap(),
-      Expression.Const.make(wasm_mod, const_int32 @@ Buffer.length(buf)),
-    ),
-  ];
-  let elts =
-    List.mapi(
-      (idx, i: int64) =>
+      gensym_label(label),
+      [
         store(
-          ~ty=Type.int64,
-          ~offset=8 * (idx + 1),
+          ~offset=0,
+          wasm_mod,
+          // compute number of words (add 3 to account for floored division)
+          tee_swap(heap_allocate(wasm_mod, env, 2 + (data_size + 3) / 4)),
+          Expression.Const.make(
+            wasm_mod,
+            const_int32(tag_val_of_heap_tag_type(tag)),
+          ),
+        ),
+        store(
+          ~offset=4,
           wasm_mod,
           get_swap(),
-          Expression.Const.make(wasm_mod, wrap_int64(i)),
+          Expression.Const.make(wasm_mod, const_int32 @@ data_size),
+        ),
+        Expression.Memory_init.make(
+          wasm_mod,
+          segment_offset,
+          Expression.Binary.make(
+            wasm_mod,
+            Op.add_int32,
+            get_swap(),
+            Expression.Const.make(wasm_mod, const_int32(8)),
+          ),
+          Expression.Const.make(wasm_mod, const_int32(0)),
+          Expression.Const.make(wasm_mod, const_int32(data_size)),
+          grain_memory,
         ),
-      ints_to_push,
+        get_swap(),
+      ],
     );
-  Expression.Block.make(
-    wasm_mod,
-    gensym_label(label),
-    List.concat([preamble, elts, [get_swap()]]),
-  );
+  } else {
+    let ints_to_push: list(int64) = buf_to_ints(buf);
+    let preamble = [
+      store(
+        ~offset=0,
+        wasm_mod,
+        tee_swap(
+          heap_allocate(wasm_mod, env, 2 + 2 * List.length(ints_to_push)),
+        ),
+        Expression.Const.make(
+          wasm_mod,
+          const_int32(tag_val_of_heap_tag_type(tag)),
+        ),
+      ),
+      store(
+        ~offset=4,
+        wasm_mod,
+        get_swap(),
+        Expression.Const.make(wasm_mod, const_int32 @@ Buffer.length(buf)),
+      ),
+    ];
+    let elts =
+      List.mapi(
+        (idx, i: int64) =>
+          store(
+            ~ty=Type.int64,
+            ~offset=8 * (idx + 1),
+            wasm_mod,
+            get_swap(),
+            Expression.Const.make(wasm_mod, wrap_int64(i)),
+          ),
+        ints_to_push,
+      );
+    Expression.Block.make(
+      wasm_mod,
+      gensym_label(label),
+      List.concat([preamble, elts, [get_swap()]]),
+    );
+  };
 };
 
 let allocate_byte_like_uninitialized = (wasm_mod, env, size, tag, label) => {
@@ -3378,7 +3438,7 @@ let compile_type_metadata = (wasm_mod, env, prog) => {
     | (initial_memory, None) => (initial_memory, Memory.unlimited)
     };
 
-  let data_segments =
+  let segments =
     switch (metadata_tbl_data) {
     | Some(data) => [
         Memory.{
@@ -3396,6 +3456,14 @@ let compile_type_metadata = (wasm_mod, env, prog) => {
       ]
     | None => []
     };
+  let data_segments =
+    if (Config.bulk_memory^) {
+      let (_, data_segments) = data_segments^;
+      List.rev_append(data_segments, segments);
+    } else {
+      segments;
+    };
+
   Memory.set_memory(
     wasm_mod,
     initial_memory,
@@ -3461,8 +3529,6 @@ let compile_wasm_module =
       Option.is_none(Grain_utils.Config.memory_base^),
     );
 
-  compile_type_metadata(wasm_mod, env, prog);
-
   ignore @@
   Table.add_table(
     wasm_mod,
@@ -3497,12 +3563,15 @@ let compile_wasm_module =
 
   ignore @@ compile_main(wasm_mod, env, prog);
 
+  compile_type_metadata(wasm_mod, env, prog);
+
   validate_module(~name?, wasm_mod);
 
   switch (Config.profile^) {
   | Some(Release) => Optimize_mod.optimize(wasm_mod)
   | None => ()
   };
+
   wasm_mod;
 };
 

compiler/src/utils/wasm_utils.re

@@ -22,7 +22,7 @@ type wasm_bin_section_type =
   | Element
   | Code
   | Data
-  | DataCount;
+  | DataCount(int);
 
 [@deriving sexp]
 type wasm_bin_section = {
@@ -213,7 +213,7 @@ let section_type_of_int = (~pos=?, ~name=?) =>
   | 9 => Element
   | 10 => Code
   | 11 => Data
-  | 12 => DataCount
+  | 12 => DataCount(-1)
   | n => raise(MalformedSectionType(n, pos));
 
 let int_of_section_type =
@@ -230,7 +230,7 @@ let int_of_section_type =
   | Element => 9
   | Code => 10
   | Data => 11
-  | DataCount => 12;
+  | DataCount(_) => 12;
 
 let get_wasm_sections = (~reset=false, inchan) => {
   let orig_pos = pos_in(inchan);

compiler/src/utils/wasm_utils.rei

@@ -21,7 +21,7 @@ type wasm_bin_section_type =
   | Element
   | Code
   | Data
-  | DataCount;
+  | DataCount(int);
 
 [@deriving sexp]
 type wasm_bin_section = {

compiler/test/suites/basic_functionality.re

@@ -377,6 +377,6 @@ describe("basic functionality", ({test, testSkip}) => {
     ~config_fn=smallestFileConfig,
     "smallest_grain_program",
     "",
-    6507,
+    6281,
   );
 });