[wasm] Option to implement OCaml strings with JavaScript strings

.github/workflows/build-wasm_of_ocaml.yml

@@ -75,7 +75,7 @@ jobs:
       - name: Set-up Node.js
         uses: actions/setup-node@v4
         with:
-          node-version: latest
+          node-version: 'v24.0.0-v8-canary202412116884e26428'
 
       - name: Set-up OCaml ${{ matrix.ocaml-compiler }}
         uses: ocaml/setup-ocaml@v3

benchmarks/Makefile

@@ -17,12 +17,12 @@ bench:
 	@date -u +"%FT%TZ - Installing dependencies"
 	opam pin -yn --with-version=dev ..
 	opam install -y wasm_of_ocaml-compiler js_of_ocaml-ppx gen_js_api brr
-	$(MAKE) microbenchmarks
-	$(MAKE) -C benchmark-fiat-crypto bench
-	$(MAKE) -C benchmark-ocamlc bench
+	#$(MAKE) microbenchmarks
+	#$(MAKE) -C benchmark-fiat-crypto bench
+	#$(MAKE) -C benchmark-ocamlc bench
 	$(MAKE) -C benchmark-partial-render-table bench
-	$(MAKE) -C benchmark-camlboy bench
-	$(MAKE) -C benchmark-others bench
+	#$(MAKE) -C benchmark-camlboy bench
+	#$(MAKE) -C benchmark-others bench # To try later!
 
 microbenchmarks:
 	@date -u +"%FT%TZ - Microbenchmarks: starting"

benchmarks/benchmark-partial-render-table/Makefile

@@ -14,7 +14,7 @@ bench:
 	@date -u +"%FT%TZ - $(NAME): done"
 
 perform:
-	/usr/bin/time -f "%E %R" $(COMPILER) --debug times --opt 2 --pretty main.bc-for-jsoo -o out.js 2>&1 | \
+	/usr/bin/time -f "%E %R" $(COMPILER) --debug times --opt 2 --pretty --enable use-js-string main.bc-for-jsoo -o out.js 2>&1 | \
 	tee /dev/stderr | \
 	ocaml -I +str str.cma ../utils/compilation_metrics.ml $(COMPILER) "$(NAME)" out.js | \
 	sh ../utils/aggregate.sh $(KIND)

compiler/bin-wasm_of_ocaml/compile.ml

@@ -84,6 +84,7 @@ let preprocessor_variables () =
         | `Disabled | `Jspi -> "jspi"
         | `Cps -> "cps"
         | `Double_translation -> assert false) )
+  ; "use-js-string", Wat_preprocess.Bool (Config.Flag.use_js_string ())
   ]
 
 let with_runtime_files ~runtime_wasm_files f =
@@ -249,25 +250,25 @@ let generate_prelude ~out_file =
     Driver.optimize_for_wasm ~profile code
   in
   let context = Generate.start () in
-  let _ =
+  let _, generated_js =
     Generate.f
       ~context
-      ~unit_name:(Some "prelude")
+      ~unit_name:(Some "wasmoo_prelude")
       ~live_vars:variable_uses
       ~in_cps
       ~deadcode_sentinal
       ~global_flow_data
       program
   in
   Generate.wasm_output ch ~opt_source_map_file:None ~context;
-  uinfo.provides
+  uinfo.provides, generated_js
 
 let build_prelude z =
   Fs.with_intermediate_file (Filename.temp_file "prelude" ".wasm")
   @@ fun prelude_file ->
-  let predefined_exceptions = generate_prelude ~out_file:prelude_file in
+  let info = generate_prelude ~out_file:prelude_file in
   Zip.add_file z ~name:"prelude.wasm" ~file:prelude_file;
-  predefined_exceptions
+  info
 
 let build_js_runtime ~primitives ?runtime_arguments () =
   let always_required_js, primitives =
@@ -443,12 +444,17 @@ let run
      let z = Zip.open_out tmp_output_file in
      Zip.add_file z ~name:"runtime.wasm" ~file:tmp_wasm_file;
      Zip.add_entry z ~name:"runtime.js" ~contents:js_runtime;
-     let predefined_exceptions = build_prelude z in
+     let predefined_exceptions, fragments = build_prelude z in
      Link.add_info
        z
        ~predefined_exceptions
        ~build_info:(Build_info.create `Runtime)
-       ~unit_data:[]
+       ~unit_data:
+         [ { Link.unit_name = "wasmoo_prelude"
+           ; unit_info = Unit_info.empty
+           ; fragments
+           }
+         ]
        ();
      Zip.close_out z)
    else

compiler/bin-wasm_of_ocaml/gen/gen.ml

@@ -27,7 +27,12 @@ let check_js_file fname =
 
 let default_flags = []
 
-let interesting_runtimes = [ [ "effects", `S "jspi" ]; [ "effects", `S "cps" ] ]
+let interesting_runtimes =
+  [ [ "effects", `S "jspi"; "use-js-string", `B false ]
+  ; [ "effects", `S "cps"; "use-js-string", `B false ]
+  ; [ "effects", `S "jspi"; "use-js-string", `B true ]
+  ; [ "effects", `S "cps"; "use-js-string", `B true ]
+  ]
 
 let name_runtime standard l =
   let flags =

compiler/lib-wasm/gc_target.ml

@@ -35,14 +35,31 @@ module Type = struct
           ; typ = W.Array { mut = true; typ = Value value }
           })
 
-  let string_type =
-    register_type "string" (fun () ->
+  let bytes_type =
+    register_type "bytes" (fun () ->
         return
           { supertype = None
           ; final = true
           ; typ = W.Array { mut = true; typ = Packed I8 }
           })
 
+  let string_type =
+    register_type "string" (fun () ->
+        return
+          (if Config.Flag.use_js_string ()
+           then
+             { supertype = None
+             ; final = true
+             ; typ =
+                 W.Struct
+                   [ { mut = false; typ = Value (Ref { nullable = true; typ = Any }) } ]
+             }
+           else
+             { supertype = None
+             ; final = true
+             ; typ = W.Array { mut = true; typ = Packed I8 }
+             }))
+
   let float_type =
     register_type "float" (fun () ->
         return
@@ -788,13 +805,48 @@ module Memory = struct
        wasm_array_set ~ty:Type.float_array_type (load a) (load i) (unbox_float (load v)))
 
   let bytes_length e =
-    let* ty = Type.string_type in
+    let* ty = Type.bytes_type in
     let* e = wasm_cast ty e in
     return (W.ArrayLen e)
 
-  let bytes_get e e' = wasm_array_get ~ty:Type.string_type e e'
+  let bytes_get e e' = wasm_array_get ~ty:Type.bytes_type e e'
+
+  let bytes_set e e' e'' = wasm_array_set ~ty:Type.bytes_type e e' e''
 
-  let bytes_set e e' e'' = wasm_array_set ~ty:Type.string_type e e' e''
+  let string_value e =
+    let* string = Type.string_type in
+    let* e = wasm_struct_get string (wasm_cast string e) 0 in
+    return (W.ExternConvertAny e)
+
+  let string_length e =
+    if Config.Flag.use_js_string ()
+    then
+      let* f =
+        register_import
+          ~import_module:"wasm:js-string"
+          ~name:"length"
+          (Fun { W.params = [ Ref { nullable = true; typ = Extern } ]; result = [ I32 ] })
+      in
+      let* e = string_value e in
+      return (W.Call (f, [ e ]))
+    else bytes_length e
+
+  let string_get e e' =
+    if Config.Flag.use_js_string ()
+    then
+      let* f =
+        register_import
+          ~import_module:"wasm:js-string"
+          ~name:"charCodeAt"
+          (Fun
+             { W.params = [ Ref { nullable = true; typ = Extern }; I32 ]
+             ; result = [ I32 ]
+             })
+      in
+      let* e = string_value e in
+      let* e' = Value.int_val e' in
+      Value.val_int (return (W.Call (f, [ e; e' ])))
+    else bytes_get e e'
 
   let field e idx = wasm_array_get e (Arith.const (Int32.of_int (idx + 1)))
 
@@ -913,6 +965,17 @@ module Constant = struct
     | Const_named of string
     | Mutated
 
+  let translate_js_string s =
+    let* x =
+      register_import
+        ~import_module:"str"
+        ~name:s
+        (Global { mut = false; typ = Ref { nullable = false; typ = Extern } })
+    in
+    let* ty = Type.js_type in
+    return
+      (Const_named ("str_" ^ s), W.StructNew (ty, [ AnyConvertExtern (GlobalGet x) ]))
+
   let rec translate_rec c =
     match c with
     | Code.Int i -> return (Const, W.RefI31 (Const (I32 (Targetint.to_int32 i))))
@@ -971,34 +1034,29 @@ module Constant = struct
           | Utf (Utf8 s) -> s
           | Byte s -> byte_string s
         in
-        let* x =
-          register_import
-            ~import_module:"str"
-            ~name:s
-            (Global { mut = false; typ = Ref { nullable = false; typ = Extern } })
-        in
-        let* ty = Type.js_type in
-        return
-          (Const_named ("str_" ^ s), W.StructNew (ty, [ AnyConvertExtern (GlobalGet x) ]))
+        translate_js_string s
     | String s ->
-        let* ty = Type.string_type in
-        if String.length s >= string_length_threshold
-        then
-          let name = Code.Var.fresh_n "string" in
-          let* () = register_data_segment name s in
-          return
-            ( Mutated
-            , W.ArrayNewData
-                (ty, name, Const (I32 0l), Const (I32 (Int32.of_int (String.length s))))
-            )
+        if Config.Flag.use_js_string ()
+        then translate_js_string (byte_string s)
         else
-          let l =
-            String.fold_right
-              ~f:(fun c r -> W.Const (I32 (Int32.of_int (Char.code c))) :: r)
-              s
-              ~init:[]
-          in
-          return (Const_named ("str_" ^ s), W.ArrayNewFixed (ty, l))
+          let* ty = Type.string_type in
+          if String.length s >= string_length_threshold
+          then
+            let name = Code.Var.fresh_n "string" in
+            let* () = register_data_segment name s in
+            return
+              ( Mutated
+              , W.ArrayNewData
+                  (ty, name, Const (I32 0l), Const (I32 (Int32.of_int (String.length s))))
+              )
+          else
+            let l =
+              String.fold_right
+                ~f:(fun c r -> W.Const (I32 (Int32.of_int (Char.code c))) :: r)
+                s
+                ~init:[]
+            in
+            return (Const_named ("str_" ^ s), W.ArrayNewFixed (ty, l))
     | Float f ->
         let* ty = Type.float_type in
         return (Const, W.StructNew (ty, [ Const (F64 (Int64.float_of_bits f)) ]))

compiler/lib-wasm/generate.ml

@@ -300,29 +300,30 @@ module Generate (Target : Target_sig.S) = struct
         seq (Memory.array_set x y z) Value.unit);
     register_tern_prim "caml_floatarray_unsafe_set" ~ty:(Int Normalized) (fun x y z ->
         seq (Memory.float_array_set x y z) Value.unit);
-    register_bin_prim "caml_string_unsafe_get" `Pure ~ty:(Int Normalized) Memory.bytes_get;
+    register_bin_prim "caml_string_unsafe_get" `Pure ~ty:(Int Normalized) Memory.string_get;
     register_bin_prim
       "caml_bytes_unsafe_get"
       `Mutable
       ~ty:(Int Normalized)
       Memory.bytes_get;
-    register_tern_prim
-      "caml_string_unsafe_set"
-      ~ty:(Int Normalized)
-      ~tz:(Int Unnormalized)
-      (fun x y z -> seq (Memory.bytes_set x y z) Value.unit);
     register_tern_prim
       "caml_bytes_unsafe_set"
       ~ty:(Int Normalized)
       ~tz:(Int Unnormalized)
       (fun x y z -> seq (Memory.bytes_set x y z) Value.unit);
+    let string_get context x y =
+      seq
+        (let* cond = Arith.uge (Value.int_val y) (Memory.string_length x) in
+         instr (W.Br_if (label_index context bound_error_pc, cond)))
+        (Memory.string_get x y)
+    in
+    register_bin_prim_ctx "caml_string_get" ~ty:(Int Normalized) string_get;
     let bytes_get context x y =
       seq
         (let* cond = Arith.uge y (Memory.bytes_length x) in
          instr (W.Br_if (label_index context bound_error_pc, cond)))
         (Memory.bytes_get x y)
     in
-    register_bin_prim_ctx "caml_string_get" ~ty:(Int Normalized) bytes_get;
     register_bin_prim_ctx "caml_bytes_get" ~ty:(Int Normalized) bytes_get;
     let bytes_set context x y z =
       seq
@@ -331,17 +332,12 @@ module Generate (Target : Target_sig.S) = struct
          Memory.bytes_set x y z)
         Value.unit
     in
-    register_tern_prim_ctx
-      "caml_string_set"
-      ~ty:(Int Normalized)
-      ~tz:(Int Unnormalized)
-      bytes_set;
     register_tern_prim_ctx
       "caml_bytes_set"
       ~ty:(Int Normalized)
       ~tz:(Int Unnormalized)
       bytes_set;
-    register_un_prim "caml_ml_string_length" `Pure (fun x -> Memory.bytes_length x);
+    register_un_prim "caml_ml_string_length" `Pure (fun x -> Memory.string_length x);
     register_un_prim "caml_ml_bytes_length" `Pure (fun x -> Memory.bytes_length x);
     register_bin_prim
       "%int_add"
@@ -1025,7 +1021,6 @@ module Generate (Target : Target_sig.S) = struct
                     ( Extern
                         ( "caml_string_get"
                         | "caml_bytes_get"
-                        | "caml_string_set"
                         | "caml_bytes_set"
                         | "caml_check_bound"
                         | "caml_check_bound_gen"

compiler/lib-wasm/link.ml

@@ -393,7 +393,7 @@ let generate_start_function ~to_link ~out_file =
   Filename.gen_file out_file
   @@ fun ch ->
   let context = Generate.start () in
-  Generate.add_init_function ~context ~to_link:("prelude" :: to_link);
+  Generate.add_init_function ~context ~to_link:("wasmoo_prelude" :: to_link);
   Generate.wasm_output ch ~opt_source_map_file:None ~context;
   if times () then Format.eprintf "    generate start: %a@." Timer.print t1
 
@@ -644,11 +644,11 @@ let load_information files =
   match files with
   | [] -> assert false
   | runtime :: other_files ->
-      let build_info, predefined_exceptions, _unit_data =
+      let build_info, predefined_exceptions, unit_data =
         Zip.with_open_in runtime read_info
       in
       ( predefined_exceptions
-      , (runtime, (build_info, []))
+      , (runtime, (build_info, unit_data))
         :: List.map other_files ~f:(fun file ->
                let build_info, _predefined_exceptions, unit_data =
                  Zip.with_open_in file read_info
@@ -745,7 +745,8 @@ let link ~output_file ~linkall ~enable_source_maps ~files =
               || cmo_file
               || linkall
               || unit_info.force_link
-              || not (StringSet.is_empty (StringSet.inter requires unit_info.provides))
+              || (not (StringSet.is_empty (StringSet.inter requires unit_info.provides)))
+              || String.equal unit_name "wasmoo_prelude"
             then
               ( StringSet.diff
                   (StringSet.union unit_info.requires requires)

compiler/lib-wasm/runtime.ml

@@ -24,7 +24,9 @@ let build ~allowed_imports ~link_options ~opt_options ~variables ~inputs ~output
       let missing_imports =
         List.filter
           ~f:(fun { Link.Wasm_binary.module_; _ } ->
-            not (List.mem ~eq:String.equal module_ allowed_imports))
+            not
+              (String.equal module_ ""
+              || List.mem ~eq:String.equal module_ allowed_imports))
           imports
       in
       if not (List.is_empty missing_imports)

compiler/lib-wasm/target_sig.ml

@@ -80,6 +80,10 @@ module type S = sig
 
     val bytes_set : expression -> expression -> expression -> unit Code_generation.t
 
+    val string_length : expression -> expression
+
+    val string_get : expression -> expression -> expression
+
     val box_float : expression -> expression
 
     val unbox_float : expression -> expression