Support for constant JavaScript strings

compiler/bin-js_of_ocaml/build_fs.ml

@@ -73,7 +73,7 @@ function jsoo_create_file_extern(name,content){
   let code = Code.prepend Code.empty instr in
   Filename.gen_file output_file (fun chan ->
       let pfs_fmt = Pretty_print.to_out_channel chan in
-      let (_ : Source_map.t option) =
+      let (_ : Source_map.t option), _ =
         Driver.f
           ~target:(`JavaScript pfs_fmt)
           ~standalone:true

compiler/bin-js_of_ocaml/compile.ml

@@ -179,7 +179,7 @@ let run
   let output (one : Parse_bytecode.one) ~standalone ~source_map ~linkall output_file =
     check_debug one;
     let init_pseudo_fs = fs_external && standalone in
-    let sm =
+    let sm, _ =
       match output_file with
       | `Stdout, fmt ->
           let instr =

compiler/bin-wasm_of_ocaml/compile.ml

@@ -162,7 +162,7 @@ let escape_string s =
   done;
   Buffer.contents b
 
-let build_js_runtime primitives wasm_file output_file =
+let build_js_runtime primitives (strings : string list) wasm_file output_file =
   let always_required_js, primitives =
     let l =
       StringSet.fold
@@ -187,15 +187,28 @@ let build_js_runtime primitives wasm_file output_file =
   let f = Pretty_print.to_buffer b' in
   Pretty_print.set_compact f (not (Config.Flag.pretty ()));
   ignore (Js_output.program f [ primitives ]);
+  let b'' = Buffer.create 1024 in
+  let f = Pretty_print.to_buffer b'' in
+  Pretty_print.set_compact f (not (Config.Flag.pretty ()));
+  ignore
+    (Js_output.program
+       f
+       [ ( Javascript.Expression_statement
+             (EArr
+                (List.map
+                   ~f:(fun s -> Javascript.Element (EStr (Utf8_string.of_string_exn s)))
+                   strings))
+         , Javascript.N )
+       ]);
   let s = Wa_runtime.js_runtime in
   let rec find pat i =
     if String.equal (String.sub s ~pos:i ~len:(String.length pat)) pat
     then i
     else find pat (i + 1)
   in
-  let i = String.index s '\n' + 1 in
-  let j = find "CODE" 0 in
-  let k = find "PRIMITIVES" 0 in
+  let i = find "CODE" 0 in
+  let j = find "PRIMITIVES" 0 in
+  let k = find "STRINGS" 0 in
   let rec trim_semi s =
     let l = String.length s in
     if l = 0
@@ -207,13 +220,14 @@ let build_js_runtime primitives wasm_file output_file =
   in
   write_file
     output_file
-    (String.sub s ~pos:0 ~len:i
-    ^ Buffer.contents b
-    ^ String.sub s ~pos:i ~len:(j - i)
+    (Buffer.contents b
+    ^ String.sub s ~pos:0 ~len:i
     ^ escape_string (Filename.basename wasm_file)
-    ^ String.sub s ~pos:(j + 4) ~len:(k - j - 4)
+    ^ String.sub s ~pos:(i + 4) ~len:(j - i - 4)
     ^ trim_semi (Buffer.contents b')
-    ^ String.sub s ~pos:(k + 10) ~len:(String.length s - k - 10))
+    ^ String.sub s ~pos:(j + 10) ~len:(k - j - 10)
+    ^ trim_semi (Buffer.contents b'')
+    ^ String.sub s ~pos:(k + 7) ~len:(String.length s - k - 7))
 
 let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; params } =
   Wa_generate.init ();
@@ -256,7 +270,7 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
   let need_debug = Config.Flag.debuginfo () in
   let output (one : Parse_bytecode.one) ~standalone ch =
     let code = one.code in
-    let _ =
+    let _, strings =
       Driver.f
         ~target:(`Wasm ch)
         ~standalone
@@ -266,7 +280,8 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
         one.debug
         code
     in
-    if times () then Format.eprintf "compilation: %a@." Timer.print t
+    if times () then Format.eprintf "compilation: %a@." Timer.print t;
+    strings
   in
   (let kind, ic, close_ic, include_dirs =
      let ch = open_in_bin input_file in
@@ -296,9 +311,9 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
        if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
        let wat_file = Filename.chop_extension (fst output_file) ^ ".wat" in
        let wasm_file = Filename.chop_extension (fst output_file) ^ ".wasm" in
-       output_gen wat_file (output code ~standalone:true);
+       let strings = output_gen wat_file (output code ~standalone:true) in
        let primitives = link_and_optimize runtime_wasm_files wat_file wasm_file in
-       build_js_runtime primitives wasm_file (fst output_file)
+       build_js_runtime primitives strings wasm_file (fst output_file)
    | `Cmo _ | `Cma _ -> assert false);
    close_ic ());
   Debug.stop_profiling ()

compiler/lib/driver.ml

@@ -613,14 +613,15 @@ let full ~target ~standalone ~wrap_with_fun ~profile ~linkall ~source_map d p =
   let r = opt p in
   let () = if times () then Format.eprintf " optimizations : %a@." Timer.print t in
   match target with
-  | `JavaScript formatter -> emit formatter r
+  | `JavaScript formatter ->
+      let source_map = emit formatter r in
+      source_map, []
   | `Wasm ch ->
       let (p, live_vars), _, in_cps = r in
-      Wa_generate.f ch ~live_vars ~in_cps p;
-      None
+      None, Wa_generate.f ch ~live_vars ~in_cps p
 
 let full_no_source_map ~target ~standalone ~wrap_with_fun ~profile ~linkall d p =
-  let (_ : Source_map.t option) =
+  let (_ : Source_map.t option * string list) =
     full ~target ~standalone ~wrap_with_fun ~profile ~linkall ~source_map:None d p
   in
   ()

compiler/lib/driver.mli

@@ -29,7 +29,7 @@ val f :
   -> ?source_map:Source_map.t
   -> Parse_bytecode.Debug.t
   -> Code.program
-  -> Source_map.t option
+  -> Source_map.t option * string list
 
 val f' :
      ?standalone:bool

compiler/lib/specialize_js.ml

@@ -123,13 +123,13 @@ let specialize_instr ~target info i =
       match the_native_string_of info f with
       | Some s -> Let (x, Prim (Extern "caml_js_delete", [ o; Pc (NativeString s) ]))
       | _ -> i)
-  | ( Let (x, Prim (Extern ("caml_jsstring_of_string" | "caml_js_from_string"), [ y ]))
-    , `JavaScript ) -> (
+  | Let (x, Prim (Extern ("caml_jsstring_of_string" | "caml_js_from_string"), [ y ])), _
+    -> (
       match the_string_of info y with
       | Some s when String.is_valid_utf_8 s ->
           Let (x, Constant (NativeString (Native_string.of_string s)))
       | Some _ | None -> i)
-  | Let (x, Prim (Extern "caml_jsbytes_of_string", [ y ])), `JavaScript -> (
+  | Let (x, Prim (Extern "caml_jsbytes_of_string", [ y ])), _ -> (
       match the_string_of info y with
       | Some s -> Let (x, Constant (NativeString (Native_string.of_bytestring s)))
       | None -> i)

compiler/lib/stdlib.ml

@@ -1163,13 +1163,16 @@ module Filename = struct
     in
     try
       let ch = open_out_bin f_tmp in
-      (try f ch
-       with e ->
-         close_out ch;
-         raise e);
+      let res =
+        try f ch
+        with e ->
+          close_out ch;
+          raise e
+      in
       close_out ch;
       (try Sys.remove file with Sys_error _ -> ());
-      Sys.rename f_tmp file
+      Sys.rename f_tmp file;
+      res
     with exc ->
       Sys.remove f_tmp;
       raise exc

compiler/lib/wasm/wa_ast.ml

@@ -11,6 +11,7 @@ type packed_type =
 type heap_type =
   | Func
   | Extern
+  | Any
   | Eq
   | I31
   | Type of var

compiler/lib/wasm/wa_code_generation.ml

@@ -34,6 +34,8 @@ type context =
   ; mutable dummy_funs : Var.t IntMap.t
   ; mutable cps_dummy_funs : Var.t IntMap.t
   ; mutable init_code : W.instruction list
+  ; mutable string_count : int
+  ; mutable strings : string list
   }
 
 let make_context () =
@@ -51,6 +53,8 @@ let make_context () =
   ; dummy_funs = IntMap.empty
   ; cps_dummy_funs = IntMap.empty
   ; init_code = []
+  ; string_count = 0
+  ; strings = []
   }
 
 type var =
@@ -171,6 +175,13 @@ let register_init_code code st =
   st.context.init_code <- st'.instrs @ st.context.init_code;
   (), st
 
+let register_string s st =
+  let context = st.context in
+  let n = context.string_count in
+  context.string_count <- 1 + context.string_count;
+  context.strings <- s :: context.strings;
+  n, st
+
 let set_closure_env f env st =
   st.context.closure_envs <- Var.Map.add f env st.context.closure_envs;
   (), st

compiler/lib/wasm/wa_code_generation.mli

@@ -18,6 +18,8 @@ type context =
   ; mutable dummy_funs : Code.Var.t Stdlib.IntMap.t
   ; mutable cps_dummy_funs : Code.Var.t Stdlib.IntMap.t
   ; mutable init_code : Wa_ast.instruction list
+  ; mutable string_count : int
+  ; mutable strings : string list
   }
 
 val make_context : unit -> context
@@ -136,6 +138,8 @@ val register_init_code : unit t -> unit t
 
 val init_code : context -> unit t
 
+val register_string : string -> int t
+
 val get_context : context t
 
 val set_closure_env : Code.Var.t -> Code.Var.t -> unit t

compiler/lib/wasm/wa_gc_target.ml

@@ -41,6 +41,16 @@ module Type = struct
           ; typ = W.Array { mut = true; typ = Value F64 }
           })
 
+  let js_type =
+    register_type "js" (fun () ->
+        return
+          { supertype = None
+          ; final = true
+          ; typ =
+              W.Struct
+                [ { mut = false; typ = Value (Ref { nullable = true; typ = Any }) } ]
+          })
+
   let compare_type =
     register_type "compare" (fun () ->
         return
@@ -764,6 +774,23 @@ module Constant = struct
     let* () = register_global (V name) { mut = false; typ = Type.value } c in
     return (W.GlobalGet (V name))
 
+  let str_js_utf8 s =
+    let b = Buffer.create (String.length s) in
+    String.iter s ~f:(function
+        | '\\' -> Buffer.add_string b "\\\\"
+        | c -> Buffer.add_char b c);
+    Buffer.contents b
+
+  let str_js_byte s =
+    let b = Buffer.create (String.length s) in
+    String.iter s ~f:(function
+        | '\\' -> Buffer.add_string b "\\\\"
+        | '\128' .. '\255' as c ->
+            Buffer.add_string b "\\x";
+            Buffer.add_char_hex b c
+        | c -> Buffer.add_char b c);
+    Buffer.contents b
+
   let rec translate_rec c =
     match c with
     | Code.Int (Regular, i) -> return (true, W.RefI31 (Const (I32 i)))
@@ -805,7 +832,22 @@ module Constant = struct
           in
           return (true, c)
         else return (true, c)
-    | NativeString (Byte s | Utf (Utf8 s)) | String s ->
+    | NativeString s ->
+        let s =
+          match s with
+          | Utf (Utf8 s) -> str_js_utf8 s
+          | Byte s -> str_js_byte s
+        in
+        let* i = register_string s in
+        let* x =
+          register_import
+            ~import_module:"strings"
+            ~name:(string_of_int i)
+            (Global { mut = false; typ = Ref { nullable = false; typ = Any } })
+        in
+        let* ty = Type.js_type in
+        return (true, W.StructNew (ty, [ GlobalGet (V x) ]))
+    | String s ->
         let* ty = Type.string_type in
         if String.length s > string_length_threshold
         then

compiler/lib/wasm/wa_generate.ml

@@ -994,9 +994,10 @@ module Generate (Target : Wa_target_sig.S) = struct
           W.Data { name; read_only = true; active; contents })
         (Var.Map.bindings ctx.global_context.data_segments)
     in
-    List.rev_append
-      ctx.global_context.other_fields
-      (imports @ functions @ (start_function :: constant_data))
+    ( List.rev_append
+        ctx.global_context.other_fields
+        (imports @ functions @ (start_function :: constant_data))
+    , List.rev ctx.global_context.strings )
 end
 
 let init () =
@@ -1047,9 +1048,11 @@ let f ch (p : Code.program) ~live_vars ~in_cps =
   match target with
   | `Core ->
       let module G = Generate (Wa_core_target) in
-      let fields = G.f ~live_vars ~in_cps p in
-      Wa_asm_output.f ch fields
+      let fields, strings = G.f ~live_vars ~in_cps p in
+      Wa_asm_output.f ch fields;
+      strings
   | `GC ->
       let module G = Generate (Wa_gc_target) in
-      let fields = G.f ~live_vars ~in_cps p in
-      Wa_wat_output.f ch fields
+      let fields, strings = G.f ~live_vars ~in_cps p in
+      Wa_wat_output.f ch fields;
+      strings

compiler/lib/wasm/wa_generate.mli

@@ -1,4 +1,8 @@
 val init : unit -> unit
 
 val f :
-  out_channel -> Code.program -> live_vars:int array -> in_cps:Effects.in_cps -> unit
+     out_channel
+  -> Code.program
+  -> live_vars:int array
+  -> in_cps:Effects.in_cps
+  -> string list

compiler/lib/wasm/wa_wat_output.ml

@@ -26,6 +26,7 @@ let heap_type (ty : heap_type) =
   match ty with
   | Func -> Atom "func"
   | Extern -> Atom "extern"
+  | Any -> Atom "any"
   | Eq -> Atom "eq"
   | I31 -> Atom "i31"
   | Type t -> index t

runtime/wasm/runtime.js

@@ -1,4 +1,4 @@
-(async function (eval_function, js) {
+(async function (eval_function, js, strings) {
     "use strict";
     const src = 'CODE';
     function loadRelative(src) {
@@ -345,7 +345,7 @@
          map_set:(m,x,v)=>m.set(x,v),
          log:(x)=>console.log('ZZZZZ', x)
         }
-    const imports = {Math:math,bindings:bindings,env:{},js:js}
+    const imports = {Math:math,bindings,env:{},js,strings}
     const wasmModule =
           isNode?await WebAssembly.instantiate(await code, imports)
                 :await WebAssembly.instantiateStreaming(code,imports)
@@ -377,4 +377,4 @@
     }
     await _initialize();
 })(((joo_global_object,jsoo_exports,globalThis)=>(x)=>eval(x))(globalThis,globalThis?.module?.exports||globalThis,globalThis),
-   PRIMITIVES);
+   PRIMITIVES, STRINGS);