Have physical equality inspect Javascript objects

Currently in wasm, physical equality behaves differently from in
Javascript: Javascript objects are boxed, and non-equality of the
pointers causes the values to be physically not equal. This can be an
issue for some folks who rely on the JS-object-inspecting semantics of
physical equality to implement things such as memoizing.

This restores the fact that physical equality inspects Javascript
objects. Other values are unaffected. This entails to pay an additional
cost of between one and two Wasm runtime type tests in physical
equality. A quick benchmark performing a few million `(==)` in an array
of integers show a slowdown of 20~25 %. Since typical programs should
perform physical equality much less frequently, we expect the overhead
to be in the noise in practice.

That being said, we may want to make this behaviour opt-in using a flag.

Co-authored-by: Jérôme Vouillon <[email protected]>

compiler/lib/wasm/wa_asm_output.ml

@@ -301,6 +301,13 @@ module Output () = struct
     | MemoryGrow (mem, e) -> expression e ^^ line (string "memory.grow " ^^ integer mem)
     | Seq (l, e') -> concat_map instruction l ^^ expression e'
     | Pop _ -> empty
+    | IfExpr (ty, e, e1, e2) ->
+        expression e
+        ^^ line (string "if" ^^ block_type { params = []; result = [ ty ] })
+        ^^ indent (expression e1)
+        ^^ line (string "else")
+        ^^ indent (expression e2)
+        ^^ line (string "end_if")
     | RefFunc _
     | Call_ref _
     | RefI31 _

compiler/lib/wasm/wa_ast.ml

@@ -157,6 +157,7 @@ type expression =
   | ExternExternalize of expression
   | Br_on_cast of int * ref_type * ref_type * expression
   | Br_on_cast_fail of int * ref_type * ref_type * expression
+  | IfExpr of value_type * expression * expression * expression
 
 and instruction =
   | Drop of expression

compiler/lib/wasm/wa_code_generation.ml

@@ -347,11 +347,13 @@ let bin_op_is_smi (op : W.int_bin_op) =
       false
   | Eq | Ne | Lt _ | Gt _ | Le _ | Ge _ -> true
 
-let is_smi e =
+let rec is_smi e =
   match e with
   | W.Const (I32 i) -> Int32.equal (Int31.wrap i) i
   | UnOp ((I32 op | I64 op), _) -> un_op_is_smi op
   | BinOp ((I32 op | I64 op), _, _) -> bin_op_is_smi op
+  | I31Get (S, _) -> true
+  | I31Get (U, _)
   | Const (I64 _ | F32 _ | F64 _)
   | ConstSym _
   | UnOp ((F32 _ | F64 _), _)
@@ -373,7 +375,6 @@ let is_smi e =
   | RefFunc _
   | Call_ref _
   | RefI31 _
-  | I31Get _
   | ArrayNew _
   | ArrayNewFixed _
   | ArrayNewData _
@@ -388,6 +389,7 @@ let is_smi e =
   | Br_on_cast _
   | Br_on_cast_fail _ -> false
   | BinOp ((F32 _ | F64 _), _, _) | RefTest _ | RefEq _ -> true
+  | IfExpr (_, _, ift, iff) -> is_smi ift && is_smi iff
 
 let get_i31_value x st =
   match st.instrs with

compiler/lib/wasm/wa_gc_target.ml

@@ -435,15 +435,72 @@ module Value = struct
 
   let le = binop Arith.( <= )
 
-  let eq i i' =
+  let ref_eq i i' =
     let* i = i in
     let* i' = i' in
-    val_int (return (W.RefEq (i, i')))
+    return (W.RefEq (i, i'))
 
-  let neq i i' =
-    let* i = i in
-    let* i' = i' in
-    val_int (Arith.eqz (return (W.RefEq (i, i'))))
+  let ref ty =
+    { W.nullable = false; typ = Type ty }
+
+  let ref_test typ e =
+    let* e = e in
+    return (W.RefTest (typ, e))
+
+  let caml_js_strict_equals x y =
+    let* x = x in
+    let* y = y in
+    let* f =
+      register_import
+        ~name:"caml_js_strict_equals"
+        ~import_module:"env"
+        (Fun { params = [ Type.value; Type.value ]; result = [ Type.value ] })
+    in
+    return (W.Call (f, [ x; y ]))
+
+  let if_expr ty cond ift iff =
+    let* cond = cond in
+    let* ift = ift in
+    let* iff = iff in
+    return (W.IfExpr (ty, cond, ift, iff))
+
+  let map f x =
+    let* x = x in
+    return (f x)
+
+  let (>>|) x f = map f x
+
+  let eq_gen ~negate x y =
+    let xv  = Code.Var.fresh () in
+    let yv  = Code.Var.fresh () in
+    let* js = Type.js_type in
+    let n =
+      if_expr
+        I32
+        (* We mimic an "and" on the two conditions, but in a way that is nicer to the
+           binaryen optimizer. *)
+        (if_expr
+          I32
+          (ref_test (ref js) (load xv))
+          (ref_test (ref js) (load yv))
+          (Arith.const 0l))
+        (caml_js_strict_equals (load xv) (load yv)
+          >>| (fun e -> W.RefCast ({ nullable = false; typ = I31 }, e))
+          >>| (fun e -> W.I31Get (S, e)))
+        (ref_eq (load xv) (load yv))
+    in
+    seq
+      (let* () = store xv x in
+        let* () = store yv y in
+        return ())
+      (val_int (if negate then Arith.eqz n else n))
+
+
+  let eq x y =
+    eq_gen ~negate:false x y
+
+  let neq x y =
+    eq_gen ~negate:true x y
 
   let ult = binop Arith.(ult)
 

compiler/lib/wasm/wa_initialize_locals.ml

@@ -51,6 +51,10 @@ let rec scan_expression ctx e =
   | Call (_, l) | ArrayNewFixed (_, l) | StructNew (_, l) -> scan_expressions ctx l
   | BlockExpr (_, l) -> scan_instructions ctx l
   | Seq (l, e') -> scan_instructions ctx (l @ [ Push e' ])
+  | IfExpr (_, cond, e1, e2) ->
+      scan_expression ctx cond;
+      scan_expression (fork_context ctx) e1;
+      scan_expression (fork_context ctx) e2
 
 and scan_expressions ctx l = List.iter ~f:(fun e -> scan_expression ctx e) l
 

compiler/lib/wasm/wa_wat_output.ml

@@ -331,6 +331,14 @@ let expression_or_instructions ctx in_function =
         ]
     | ExternInternalize e -> [ List (Atom "extern.internalize" :: expression e) ]
     | ExternExternalize e -> [ List (Atom "extern.externalize" :: expression e) ]
+    | IfExpr (ty, cond, ift, iff) ->
+        [ List
+            (Atom "if"
+            :: (block_type { params = []; result = [ ty ] })
+               @ expression cond
+               @ [ List (Atom "then" :: expression ift) ]
+               @ [ List (Atom "else" :: expression iff) ])
+        ]
   and instruction i =
     match i with
     | Drop e -> [ List (Atom "drop" :: expression e) ]