feat(compiler): Use data segments for string data

ospencer · ospencer · commit c349cc4e6c29 · 2023-09-28T08:22:56.000-05:00
diff --git a/compiler/src/codegen/compcore.re b/compiler/src/codegen/compcore.re
@@ -28,6 +28,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|];
@@ -282,6 +293,7 @@ let init_codegen_env = name => {
 
 let reset = () => {
   reset_labels();
+  reset_data_segments();
   List.iter(
     imp => imp.mimp_used = imp.mimp_mod == grain_env_mod,
     runtime_imports,
@@ -1456,45 +1468,92 @@ 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),
         ),
-      ints_to_push,
+        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)),
+        ),
+        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) => {
@@ -3514,20 +3573,7 @@ let compile_wasm_module = (~env=?, ~name=?, prog) => {
   // This is because in many use cases in which this is specified (e.g. wasm4), users
   // will expect the static region of memory below the heap base to all be available.
   let _ =
-    Settings.set_low_memory_unused(
-      Option.is_none(Grain_utils.Config.memory_base^),
-    );
-  let _ =
-    Memory.set_memory(
-      wasm_mod,
-      0,
-      Memory.unlimited,
-      "memory",
-      [],
-      false,
-      false,
-      grain_memory,
-    );
+    Settings.set_low_memory_unused(Option.is_none(Config.memory_base^));
 
   let compile_all = () => {
     ignore @@ compile_globals(wasm_mod, env, prog);
@@ -3546,6 +3592,25 @@ let compile_wasm_module = (~env=?, ~name=?, prog) => {
     compile_all();
   };
 
+  let segments =
+    if (Config.bulk_memory^) {
+      let (_, segments) = data_segments^;
+      List.rev(segments);
+    } else {
+      [];
+    };
+  let _ =
+    Memory.set_memory(
+      wasm_mod,
+      0,
+      Memory.unlimited,
+      "memory",
+      segments,
+      false,
+      false,
+      grain_memory,
+    );
+
   if (compiling_wasi_polyfill(name)) {
     write_universal_exports(
       wasm_mod,
diff --git a/compiler/src/linking/link.re b/compiler/src/linking/link.re
@@ -121,9 +121,15 @@ let is_function_imported = func =>
 //       constructing the AST (either through constructing two separate instances or by
 //       using Expression.copy())
 
-let rec globalize_names = (~function_names, ~global_names, ~label_names, expr) => {
+let rec globalize_names =
+        (~function_names, ~global_names, ~label_names, ~data_offset, expr) => {
   let globalize_names =
-    globalize_names(~function_names, ~global_names, ~label_names);
+    globalize_names(
+      ~function_names,
+      ~global_names,
+      ~label_names,
+      ~data_offset,
+    );
 
   let add_label = Hashtbl.add(label_names);
 
@@ -245,6 +251,19 @@ let rec globalize_names = (~function_names, ~global_names, ~label_names, expr) =
       globalize_names(Expression.Tuple_make.get_operand_at(expr, i));
     };
   | TupleExtract => globalize_names(Expression.Tuple_extract.get_tuple(expr))
+  | MemoryInit =>
+    Expression.Memory_init.set_segment(
+      expr,
+      Expression.Memory_init.get_segment(expr) + data_offset,
+    );
+    globalize_names(Expression.Memory_init.get_dest(expr));
+    globalize_names(Expression.Memory_init.get_offset(expr));
+    globalize_names(Expression.Memory_init.get_size(expr));
+  | DataDrop =>
+    Expression.Data_drop.set_segment(
+      expr,
+      Expression.Data_drop.get_segment(expr) + data_offset,
+    )
   | AtomicRMW
   | AtomicCmpxchg
   | AtomicWait
@@ -257,8 +276,6 @@ let rec globalize_names = (~function_names, ~global_names, ~label_names, expr) =
   | SIMDShift
   | SIMDLoad
   | SIMDLoadStoreLane
-  | MemoryInit
-  | DataDrop
   | Pop
   | RefNull
   | RefIs
@@ -387,11 +404,15 @@ let table_offset = ref(0);
 let module_id = ref(Comp_utils.encoded_int32(0));
 
 let round_to_8 = n => Int.logand(n + 7, Int.lognot(7));
+let data_offset = ref(0);
+let data_segments = ref([]);
 
 let link_all = (linked_mod, dependencies, signature) => {
   gensym_counter := 0;
   table_offset := 0;
   module_id := Comp_utils.encoded_int32(0);
+  data_offset := 0;
+  data_segments := [];
 
   let main = Module_resolution.current_filename^();
   let has_wasi_polyfill = Option.is_some(Config.wasi_polyfill^);
@@ -529,7 +550,14 @@ let link_all = (linked_mod, dependencies, signature) => {
         let mut = Global.is_mutable(global);
         let init = Global.get_init_expr(global);
 
-        globalize_names(~function_names, ~global_names, ~label_names, init);
+        let data_offset = data_offset^;
+        globalize_names(
+          ~function_names,
+          ~global_names,
+          ~label_names,
+          ~data_offset,
+          init,
+        );
         ignore @@ Global.add_global(linked_mod, new_name, ty, mut, init);
       };
     };
@@ -622,7 +650,14 @@ let link_all = (linked_mod, dependencies, signature) => {
         let num_locals = Function.get_num_vars(func);
         let locals = Array.init(num_locals, i => Function.get_var(func, i));
         let body = Function.get_body(func);
-        globalize_names(~function_names, ~global_names, ~label_names, body);
+        let data_offset = data_offset^;
+        globalize_names(
+          ~function_names,
+          ~global_names,
+          ~label_names,
+          ~data_offset,
+          body,
+        );
         ignore @@
         Function.add_function(
           linked_mod,
@@ -663,6 +698,29 @@ let link_all = (linked_mod, dependencies, signature) => {
       );
       table_offset := table_offset^ + size;
     };
+
+    let num_data_segments = Memory.get_num_segments(wasm_mod);
+    for (i in 0 to num_data_segments - 1) {
+      open Memory;
+      let data = get_segment_data(wasm_mod, i);
+      let kind =
+        if (get_segment_passive(wasm_mod, i)) {
+          Passive;
+        } else {
+          Active({
+            offset:
+              Expression.Const.make(
+                wasm_mod,
+                Literal.int32(
+                  Int32.of_int(get_segment_byte_offset(wasm_mod, i)),
+                ),
+              ),
+          });
+        };
+      let size = Bytes.length(data);
+      data_segments := [{data, kind, size}, ...data_segments^];
+    };
+    data_offset := data_offset^ + num_data_segments;
   };
 
   ignore @@
@@ -722,15 +780,16 @@ let link_all = (linked_mod, dependencies, signature) => {
             }),
           size: Bytes.length(data),
         },
+        ...data_segments^,
       ]
-    | None => []
+    | None => data_segments^
     };
   Memory.set_memory(
     linked_mod,
     initial_memory,
     maximum_memory,
     "memory",
-    data_segments,
+    List.rev(data_segments),
     false,
     false,
     Comp_utils.grain_memory,
diff --git a/compiler/src/utils/wasm_utils.re b/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);
diff --git a/compiler/src/utils/wasm_utils.rei b/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 = {
