Wasm_of_ocaml: implement use-js-string flag

Author: vouillon

compiler/bin-wasm_of_ocaml/compile.ml

@@ -90,7 +90,7 @@ let with_runtime_files ~runtime_wasm_files f =
 
 let build_runtime ~runtime_file =
   (* Keep this variables in sync with gen/gen.ml *)
-  let variables = [] in
+  let variables = [ "use-js-string", Config.Flag.use_js_string () ] in
   match
     List.find_opt Runtime_files.precompiled_runtimes ~f:(fun (flags, _) ->
         assert (

compiler/bin-wasm_of_ocaml/gen/gen.ml

@@ -5,7 +5,7 @@ let read_file ic = really_input_string ic (in_channel_length ic)
 
 let default_flags = []
 
-let interesting_runtimes = [ [] ]
+let interesting_runtimes = [ [ "use-js-string", false ]; [ "use-js-string", true ] ]
 
 let name_runtime l =
   let flags = List.filter_map (fun (k, v) -> if v then Some k else None) l in

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
@@ -794,15 +811,50 @@ 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' =
-    Value.val_int (wasm_array_get ~ty:Type.string_type e (Value.int_val e'))
+    Value.val_int (wasm_array_get ~ty:Type.bytes_type e (Value.int_val e'))
 
   let bytes_set e e' e'' =
-    wasm_array_set ~ty:Type.string_type e (Value.int_val e') (Value.int_val e'')
+    wasm_array_set ~ty:Type.bytes_type e (Value.int_val e') (Value.int_val 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)))
 
@@ -929,6 +981,21 @@ module Constant = struct
     | Const_named of string
     | Mutated
 
+  let translate_js_string s =
+    let* i = register_string s in
+    let* x =
+      let* name = unit_name in
+      register_import
+        ~import_module:
+          (match name with
+          | None -> "strings"
+          | Some name -> name ^ ".strings")
+        ~name:(string_of_int i)
+        (Global { mut = false; typ = Ref { nullable = false; typ = Any } })
+    in
+    let* ty = Type.js_type in
+    return (Const_named ("str_" ^ s), W.StructNew (ty, [ GlobalGet x ]))
+
   let rec translate_rec c =
     match c with
     | Code.Int i -> return (Const, W.RefI31 (Const (I32 (Targetint.to_int32 i))))
@@ -987,38 +1054,29 @@ module Constant = struct
           | Utf (Utf8 s) -> str_js_utf8 s
           | Byte s -> str_js_byte s
         in
-        let* i = register_string s in
-        let* x =
-          let* name = unit_name in
-          register_import
-            ~import_module:
-              (match name with
-              | None -> "strings"
-              | Some name -> name ^ ".strings")
-            ~name:(string_of_int i)
-            (Global { mut = false; typ = Ref { nullable = false; typ = Any } })
-        in
-        let* ty = Type.js_type in
-        return (Const_named ("str_" ^ s), W.StructNew (ty, [ 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 (str_js_byte 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 f) ]))

compiler/lib-wasm/generate.ml

@@ -291,22 +291,29 @@ module Generate (Target : Target_sig.S) = struct
                 seq (Memory.array_set x y z) Value.unit
             | Extern "caml_floatarray_unsafe_set", [ x; y; z ] ->
                 seq (Memory.float_array_set x y z) Value.unit
-            | Extern ("caml_string_unsafe_get" | "caml_bytes_unsafe_get"), [ x; y ] ->
-                Memory.bytes_get x y
-            | Extern ("caml_string_unsafe_set" | "caml_bytes_unsafe_set"), [ x; y; z ] ->
+            | Extern "caml_string_unsafe_get", [ x; y ] -> Memory.string_get x y
+            | Extern "caml_bytes_unsafe_get", [ x; y ] -> Memory.bytes_get x y
+            | Extern "caml_bytes_unsafe_set", [ x; y; z ] ->
                 seq (Memory.bytes_set x y z) Value.unit
-            | Extern ("caml_string_get" | "caml_bytes_get"), [ x; y ] ->
+            | Extern "caml_string_get", [ 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)
+            | Extern "caml_bytes_get", [ x; y ] ->
                 seq
                   (let* cond = Arith.uge (Value.int_val y) (Memory.bytes_length x) in
                    instr (W.Br_if (label_index context bound_error_pc, cond)))
                   (Memory.bytes_get x y)
-            | Extern ("caml_string_set" | "caml_bytes_set"), [ x; y; z ] ->
+            | Extern "caml_bytes_set", [ x; y; z ] ->
                 seq
                   (let* cond = Arith.uge (Value.int_val y) (Memory.bytes_length x) in
                    let* () = instr (W.Br_if (label_index context bound_error_pc, cond)) in
                    Memory.bytes_set x y z)
                   Value.unit
-            | Extern ("caml_ml_string_length" | "caml_ml_bytes_length"), [ x ] ->
+            | Extern "caml_ml_string_length", [ x ] ->
+                Value.val_int (Memory.string_length x)
+            | Extern "caml_ml_bytes_length", [ x ] ->
                 Value.val_int (Memory.bytes_length x)
             | Extern "%int_add", [ x; y ] -> Value.int_add x y
             | Extern "%int_sub", [ x; y ] -> Value.int_sub x y
@@ -782,7 +789,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/target_sig.ml

@@ -77,6 +77,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

runtime/wasm/array.wat

@@ -21,6 +21,7 @@
 
    (type $block (array (mut (ref eq))))
    (type $bytes (array (mut i8)))
+   (type $string (struct (field anyref)))
    (type $float (struct (field f64)))
    (type $float_array (array (mut f64)))
 

runtime/wasm/backtrace.wat

@@ -21,6 +21,7 @@
 
    (type $block (array (mut (ref eq))))
    (type $bytes (array (mut i8)))
+   (type $string (struct (field anyref)))
 
    (func (export "caml_get_exception_raw_backtrace")
       (param (ref eq)) (result (ref eq))

runtime/wasm/bigarray.wat

@@ -77,6 +77,12 @@
       (func $ta_blit_to_bytes
          (param (ref extern)) (param i32) (param (ref $bytes)) (param i32)
          (param i32)))
+   (import "bindings" "ta_blit_from_string"
+      (func $ta_blit_from_string
+         (param anyref) (param i32) (param (ref extern)) (param i32)
+         (param i32)))
+   (import "bindings" "ta_to_string"
+      (func $ta_to_string (param (ref extern)) (result (ref any))))
    (import "fail" "caml_bound_error" (func $caml_bound_error))
    (import "fail" "caml_raise_out_of_memory" (func $caml_raise_out_of_memory))
    (import "fail" "caml_invalid_argument"
@@ -126,9 +132,12 @@
       (func $caml_deserialize_int_4 (param (ref eq)) (result i32)))
    (import "marshal" "caml_deserialize_int_8"
       (func $caml_deserialize_int_8 (param (ref eq)) (result i64)))
+   (import "jsstring" "jsstring_length"
+      (func $jsstring_length (param anyref) (result i32)))
 
    (type $block (array (mut (ref eq))))
    (type $bytes (array (mut i8)))
+   (type $string (struct (field anyref)))
    (type $float (struct (field f64)))
    (type $float_array (array (mut f64)))
 
@@ -143,7 +152,13 @@
    (type $dup (func (param (ref eq)) (result (ref eq))))
    (type $custom_operations
       (struct
+(@if use-js-string
+(@then
+         (field $id (ref $string))
+)
+(@else
          (field $id (ref $bytes))
+))
          (field $compare (ref null $compare))
          (field $compare_ext (ref null $compare))
          (field $hash (ref null $hash))
@@ -2121,10 +2136,23 @@
          (i32.wrap_i64 (i64.shr_u (local.get $d) (i64.const 32))))
       (ref.i31 (i32.const 0)))
 
-   (export "caml_bytes_of_uint8_array" (func $caml_string_of_uint8_array))
-   (func $caml_string_of_uint8_array (export "caml_string_of_uint8_array")
+(@if use-js-string
+(@then
+   (func (export "caml_string_of_uint8_array")
       (param (ref eq)) (result (ref eq))
       ;; used to convert a typed array to a string
+      (local $a (ref extern))
+      (local.set $a
+         (ref.as_non_null (extern.convert_any (call $unwrap (local.get 0)))))
+      (struct.new $string (call $ta_to_string (local.get $a))))
+)
+(@else
+   (export "caml_string_of_uint8_array" (func $caml_bytes_of_uint8_array))
+))
+
+   (func $caml_bytes_of_uint8_array (export "caml_bytes_of_uint8_array")
+      (param (ref eq)) (result (ref eq))
+      ;; used to convert a typed array to bytes
       (local $a (ref extern)) (local $len i32)
       (local $s (ref $bytes))
       (local.set $a
@@ -2136,8 +2164,30 @@
          (local.get $len))
       (local.get $s))
 
-   (export "caml_uint8_array_of_bytes" (func $caml_uint8_array_of_string))
-   (func $caml_uint8_array_of_string (export "caml_uint8_array_of_string")
+(@if use-js-string
+(@then
+   (func (export "caml_uint8_array_of_string")
+      (param (ref eq)) (result (ref eq))
+      ;; Convert a string to a typed array
+      (local $ta (ref extern)) (local $len i32)
+      (local $s anyref)
+      (local.set $s
+         (struct.get $string 0 (ref.cast (ref $string) (local.get 0))))
+      (local.set $len (call $jsstring_length (local.get $s)))
+      (local.set $ta
+         (call $ta_create
+            (i32.const 3) ;; Uint8Array
+            (local.get $len)))
+      (call $ta_blit_from_string
+         (local.get $s) (i32.const 0) (local.get $ta) (i32.const 0)
+         (local.get $len))
+      (call $wrap (any.convert_extern (local.get $ta))))
+)
+(@else
+   (export "caml_uint8_array_of_string" (func $caml_uint8_array_of_bytes))
+))
+
+   (func $caml_uint8_array_of_bytes (export "caml_uint8_array_of_bytes")
       (param (ref eq)) (result (ref eq))
       ;; Convert bytes to a typed array
       (local $ta (ref extern)) (local $len i32)