add optimization of comparisons

mununki · mununki · commit 6d8439aa086f · 2024-03-13T02:29:02.000+09:00
diff --git a/jscomp/core/js_exp_make.ml b/jscomp/core/js_exp_make.ml
@@ -787,6 +787,7 @@ let rec float_equal ?comment (e0 : t) (e1 : t) : t =
       float_equal ?comment a e1
   | Number (Float { f = f0; _ }), Number (Float { f = f1 }) when f0 = f1 ->
       true_
+  | Number (Bigint { i = i0 }), Number (Bigint { i = i1 }) -> bool (i0 = i1)
   | _ -> { expression_desc = Bin (EqEqEq, e0, e1); comment }
 
 let int_equal = float_equal
@@ -1269,6 +1270,16 @@ let bigint_div ?comment (e1: t) (e2: t) = bin ?comment Div e1 e2
 
 let bigint_mod ?comment (e1: t) (e2: t) = bin ?comment Mod e1 e2
 
+let bigint_comp (cmp : Lam_compat.comparison) ?comment (e0: t) (e1: t) =
+  match (cmp, e0.expression_desc, e1.expression_desc) with
+  | Ceq, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 = i1)
+  | Cneq, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 <> i1)
+  | Cge, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 >= i1)
+  | Cgt, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 > i1)
+  | Cle, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 <= i1)
+  | Clt, Number (Bigint {i = i0; _}), Number (Bigint {i = i1; _}) -> bool (i0 < i1)
+  | _ -> bin ?comment (Lam_compile_util.jsop_of_comp cmp) e0 e1
+
 (* TODO -- alpha conversion
     remember to add parens..
 *)
diff --git a/jscomp/core/js_exp_make.mli b/jscomp/core/js_exp_make.mli
@@ -286,6 +286,8 @@ val bigint_div : ?comment: string -> t -> t -> t
 
 val bigint_mod : ?comment: string -> t -> t -> t
 
+val bigint_comp : Lam_compat.comparison -> ?comment:string -> t -> t -> t
+
 val js_comp : Lam_compat.comparison -> ?comment:string -> t -> t -> t
 
 val not : t -> t
diff --git a/jscomp/core/lam_analysis.ml b/jscomp/core/lam_analysis.ml
@@ -76,7 +76,7 @@ let rec no_side_effects (lam : Lam.t) : bool =
       (* Float operations *)
       | Pintoffloat | Pfloatofint | Pnegfloat
       (* | Pabsfloat *)
-      | Paddfloat | Psubfloat | Pmulfloat | Pdivfloat | Pfloatcomp _ | Pjscomp _
+      | Paddfloat | Psubfloat | Pmulfloat | Pdivfloat | Pfloatcomp _ | Pbigintcomp _ | Pjscomp _
       | Pnegbigint | Paddbigint | Psubbigint | Pmulbigint
       (* String operations *)
       | Pstringlength | Pstringrefu | Pstringrefs | Pbyteslength | Pbytesrefu
diff --git a/jscomp/core/lam_compile_primitive.ml b/jscomp/core/lam_compile_primitive.ml
@@ -251,6 +251,8 @@ let translate output_prefix loc (cxt : Lam_compile_context.t)
          [Not_found] or [Invalid_argument] ?
       *)
       match args with [ e1; e2 ] -> E.int_comp cmp e1 e2 | _ -> assert false)
+  | Pbigintcomp cmp -> (
+      match args with [ e1; e2 ] -> E.bigint_comp cmp e1 e2 | _ -> assert false)
   (* List --> stamp = 0
      Assert_false --> stamp = 26
   *)
diff --git a/jscomp/core/lam_convert.ml b/jscomp/core/lam_convert.ml
@@ -260,6 +260,7 @@ let lam_prim ~primitive:(p : Lambda.primitive) ~args loc : Lam.t =
   | Pmulbigint -> prim ~primitive:Pmulbigint ~args loc
   | Pdivbigint _is_safe (*FIXME*) -> prim ~primitive:Pdivbigint ~args loc
   | Pmodbigint _is_safe (*FIXME*) -> prim ~primitive:Pmodbigint ~args loc
+  | Pbigintcomp x -> prim ~primitive:(Pbigintcomp x) ~args loc
   | Pintcomp x -> prim ~primitive:(Pintcomp x) ~args loc
   | Poffsetint x -> prim ~primitive:(Poffsetint x) ~args loc
   | Poffsetref x -> prim ~primitive:(Poffsetref x) ~args loc
diff --git a/jscomp/core/lam_dispatch_primitive.ml b/jscomp/core/lam_dispatch_primitive.ml
@@ -123,6 +123,9 @@ let translate loc (prim_name : string) (args : J.expression list) : J.expression
   | "caml_float_equal_null" | "caml_float_equal_nullable"
   | "caml_float_equal_undefined" -> (
       match args with [ e0; e1 ] -> E.float_comp Ceq e0 e1 | _ -> assert false)
+  | "caml_bigint_equal_null" | "caml_bigint_equal_nullable"
+  | "caml_bigint_equal_undefined" -> (
+      match args with [ e0; e1 ] -> E.bigint_comp Ceq e0 e1 | _ -> assert false)
   | "caml_string_equal_null" | "caml_string_equal_nullable"
   | "caml_string_equal_undefined" -> (
       match args with
@@ -137,8 +140,9 @@ let translate loc (prim_name : string) (args : J.expression list) : J.expression
   | "caml_int_compare" ->
       E.runtime_call Js_runtime_modules.caml_primitive "int_compare" args
   | "caml_float_compare" -> call Js_runtime_modules.caml_primitive
+  | "caml_bigint_compare" -> call Js_runtime_modules.caml_primitive
   | "caml_string_compare" -> call Js_runtime_modules.caml_primitive
-  | "caml_bool_min" | "caml_int_min" | "caml_float_min" | "caml_string_min" -> (
+  | "caml_bool_min" | "caml_int_min" | "caml_float_min" | "caml_bigint_min" | "caml_string_min" -> (
       match args with
       | [ a; b ] ->
           if
@@ -147,7 +151,7 @@ let translate loc (prim_name : string) (args : J.expression list) : J.expression
           then E.econd (E.js_comp Clt a b) a b
           else call Js_runtime_modules.caml_primitive
       | _ -> assert false)
-  | "caml_bool_max" | "caml_int_max" | "caml_float_max" | "caml_string_max" -> (
+  | "caml_bool_max" | "caml_int_max" | "caml_float_max" | "caml_bigint_max" | "caml_string_max" -> (
       match args with
       | [ a; b ] ->
           if
diff --git a/jscomp/core/lam_primitive.ml b/jscomp/core/lam_primitive.ml
@@ -90,6 +90,7 @@ type t =
   | Pfloatcomp of Lam_compat.comparison
   | Pjscomp of Lam_compat.comparison
   | Pint64comp of Lam_compat.comparison
+  | Pbigintcomp of Lam_compat.comparison
   | Pjs_apply (*[f;arg0;arg1; arg2; ... argN]*)
   | Pjs_runtime_apply (* [f; [...]] *)
   (* String operations *)
@@ -281,6 +282,11 @@ let eq_primitive_approx (lhs : t) (rhs : t) =
       | Pfloatcomp comparison1 ->
           Lam_compat.eq_comparison comparison comparison1
       | _ -> false)
+  | Pbigintcomp comparison -> (
+      match rhs with
+      | Pbigintcomp comparison1 ->
+          Lam_compat.eq_comparison comparison comparison1
+      | _ -> false)
   | Pjscomp comparison -> (
       match rhs with
       | Pjscomp comparison1 -> Lam_compat.eq_comparison comparison comparison1
diff --git a/jscomp/core/lam_primitive.mli b/jscomp/core/lam_primitive.mli
@@ -80,6 +80,7 @@ type t =
   | Pfloatcomp of Lam_compat.comparison
   | Pjscomp of Lam_compat.comparison
   | Pint64comp of Lam_compat.comparison
+  | Pbigintcomp of Lam_compat.comparison
   | Pjs_apply (*[f;arg0;arg1; arg2; ... argN]*)
   | Pjs_runtime_apply (* [f; [...]] *)
   | Pstringlength
diff --git a/jscomp/core/lam_print.ml b/jscomp/core/lam_print.ml
@@ -140,6 +140,12 @@ let primitive ppf (prim : Lam_primitive.t) =
   | Pmulbigint -> fprintf ppf "*n"
   | Pdivbigint -> fprintf ppf "/n"
   | Pmodbigint -> fprintf ppf "modn"
+  | Pbigintcomp Ceq -> fprintf ppf "==,"
+  | Pbigintcomp Cneq -> fprintf ppf "!=,"
+  | Pbigintcomp Clt -> fprintf ppf "<,"
+  | Pbigintcomp Cle -> fprintf ppf "<=,"
+  | Pbigintcomp Cgt -> fprintf ppf ">,"
+  | Pbigintcomp Cge -> fprintf ppf ">=,"
   | Pjscomp Ceq -> fprintf ppf "#=="
   | Pjscomp Cneq -> fprintf ppf "#!="
   | Pjscomp Clt -> fprintf ppf "#<"
diff --git a/jscomp/ml/lambda.ml b/jscomp/ml/lambda.ml
@@ -230,7 +230,9 @@ type primitive =
   | Paddfloat | Psubfloat | Pmulfloat | Pdivfloat
   | Pfloatcomp of comparison
   (* Bigint operations *)
-  | Pnegbigint | Paddbigint | Psubbigint | Pmulbigint | Pdivbigint of is_safe | Pmodbigint of is_safe
+  | Pnegbigint | Paddbigint | Psubbigint
+  | Pmulbigint | Pdivbigint of is_safe | Pmodbigint of is_safe
+  | Pbigintcomp of comparison
   (* String operations *)
   | Pstringlength | Pstringrefu  | Pstringrefs
   | Pbyteslength | Pbytesrefu | Pbytessetu | Pbytesrefs | Pbytessets
diff --git a/jscomp/ml/lambda.mli b/jscomp/ml/lambda.mli
@@ -196,7 +196,9 @@ type primitive =
   | Paddfloat | Psubfloat | Pmulfloat | Pdivfloat
   | Pfloatcomp of comparison
   (* Bigint operations *)
-  | Pnegbigint | Paddbigint | Psubbigint | Pmulbigint | Pdivbigint of is_safe | Pmodbigint of is_safe
+  | Pnegbigint | Paddbigint | Psubbigint
+  | Pmulbigint | Pdivbigint of is_safe | Pmodbigint of is_safe
+  | Pbigintcomp of comparison
   (* String operations *)
   | Pstringlength | Pstringrefu  | Pstringrefs
   | Pbyteslength | Pbytesrefu | Pbytessetu | Pbytesrefs | Pbytessets
diff --git a/jscomp/ml/printlambda.ml b/jscomp/ml/printlambda.ml
@@ -185,6 +185,12 @@ let primitive ppf = function
   | Pdivbigint Unsafe -> fprintf ppf "/nu"
   | Pmodbigint Safe -> fprintf ppf "mod"
   | Pmodbigint Unsafe -> fprintf ppf "mod_unsafe"
+  | Pbigintcomp(Ceq) -> fprintf ppf "==,"
+  | Pbigintcomp(Cneq) -> fprintf ppf "!=,"
+  | Pbigintcomp(Clt) -> fprintf ppf "<,"
+  | Pbigintcomp(Cle) -> fprintf ppf "<=,"
+  | Pbigintcomp(Cgt) -> fprintf ppf ">,"
+  | Pbigintcomp(Cge) -> fprintf ppf ">=,"
   | Pstringlength -> fprintf ppf "string.length"
   | Pstringrefu -> fprintf ppf "string.unsafe_get"
   | Pstringrefs -> fprintf ppf "string.get"
@@ -292,6 +298,7 @@ let name_of_primitive = function
   | Pmulbigint -> "Pmulbigint"
   | Pdivbigint _ -> "Pdivbigint"
   | Pmodbigint _ -> "Pmodbigint"
+  | Pbigintcomp _ -> "Pbigintcomp"
   | Pstringlength -> "Pstringlength"
   | Pstringrefu -> "Pstringrefu"
   | Pstringrefs -> "Pstringrefs"
diff --git a/jscomp/ml/translcore.ml b/jscomp/ml/translcore.ml
@@ -57,6 +57,7 @@ type specialized = {
   stringcomp : Lambda.primitive;
   bytescomp : Lambda.primitive;
   int64comp : Lambda.primitive;
+  bigintcomp : Lambda.primitive;
   simplify_constant_constructor : bool;
 }
 
@@ -82,6 +83,7 @@ let comparisons_table =
             Pccall
               (Primitive.simple ~name:"caml_bytes_equal" ~arity:2 ~alloc:false);
           int64comp = Pbintcomp (Pint64, Ceq);
+          bigintcomp = Pbigintcomp Ceq;
           simplify_constant_constructor = true;
         } );
       ( "%notequal",
@@ -102,6 +104,7 @@ let comparisons_table =
               (Primitive.simple ~name:"caml_bytes_notequal" ~arity:2
                  ~alloc:false);
           int64comp = Pbintcomp (Pint64, Cneq);
+          bigintcomp = Pbigintcomp Cneq;
           simplify_constant_constructor = true;
         } );
       ( "%lessthan",
@@ -122,6 +125,7 @@ let comparisons_table =
               (Primitive.simple ~name:"caml_bytes_lessthan" ~arity:2
                  ~alloc:false);
           int64comp = Pbintcomp (Pint64, Clt);
+          bigintcomp = Pbigintcomp Clt;
           simplify_constant_constructor = false;
         } );
       ( "%greaterthan",
@@ -144,6 +148,7 @@ let comparisons_table =
               (Primitive.simple ~name:"caml_bytes_greaterthan" ~arity:2
                  ~alloc:false);
           int64comp = Pbintcomp (Pint64, Cgt);
+          bigintcomp = Pbigintcomp Cgt;
           simplify_constant_constructor = false;
         } );
       ( "%lessequal",
@@ -166,6 +171,7 @@ let comparisons_table =
               (Primitive.simple ~name:"caml_bytes_lessequal" ~arity:2
                  ~alloc:false);
           int64comp = Pbintcomp (Pint64, Cle);
+          bigintcomp = Pbigintcomp Cle;
           simplify_constant_constructor = false;
         } );
       ( "%greaterequal",
@@ -188,6 +194,7 @@ let comparisons_table =
               (Primitive.simple ~name:"caml_bytes_greaterequal" ~arity:2
                  ~alloc:false);
           int64comp = Pbintcomp (Pint64, Cge);
+          bigintcomp = Pbigintcomp Cge;
           simplify_constant_constructor = false;
         } );
       ( "%compare",
@@ -214,6 +221,9 @@ let comparisons_table =
           int64comp =
             Pccall
               (Primitive.simple ~name:"caml_int64_compare" ~arity:2 ~alloc:false);
+          bigintcomp =
+            Pccall
+              (Primitive.simple ~name:"caml_bigint_compare" ~arity:2 ~alloc:false);
           simplify_constant_constructor = false;
         } );
       ( "%bs_max",
@@ -226,6 +236,7 @@ let comparisons_table =
           floatcomp = arity2 "caml_float_max";
           stringcomp = arity2 "caml_string_max";
           int64comp = arity2 "caml_int64_max";
+          bigintcomp = arity2 "caml_bigint_max";
           simplify_constant_constructor = false;
         } );
       ( "%bs_min",
@@ -237,6 +248,7 @@ let comparisons_table =
           floatcomp = arity2 "caml_float_min";
           stringcomp = arity2 "caml_string_min";
           int64comp = arity2 "caml_int64_min";
+          bigintcomp = arity2 "caml_bigint_min";
           simplify_constant_constructor = false;
         } );
       ( "%bs_equal_null",
@@ -249,6 +261,7 @@ let comparisons_table =
           floatcomp = arity2 "caml_float_equal_null";
           stringcomp = arity2 "caml_string_equal_null";
           int64comp = arity2 "caml_int64_equal_null";
+          bigintcomp = arity2 "caml_bigint_equal_null";
           simplify_constant_constructor = true;
         } );
       ( "%bs_equal_undefined",
@@ -261,6 +274,7 @@ let comparisons_table =
           floatcomp = arity2 "caml_float_equal_undefined";
           stringcomp = arity2 "caml_string_equal_undefined";
           int64comp = arity2 "caml_int64_equal_undefined";
+          bigintcomp = arity2 "caml_bigint_equal_undefined";
           simplify_constant_constructor = true;
         } );
       ( "%bs_equal_nullable",
@@ -273,6 +287,7 @@ let comparisons_table =
           floatcomp = arity2 "caml_float_equal_nullable";
           stringcomp = arity2 "caml_string_equal_nullable";
           int64comp = arity2 "caml_int64_equal_nullable";
+          bigintcomp = arity2 "caml_bigint_equal_nullable";
           simplify_constant_constructor = true;
         } );
     ]
@@ -356,6 +371,12 @@ let primitives_table =
       ("%mulbigint", Pmulbigint);
       ("%divbigint", Pdivbigint Safe);
       ("%modbigint", Pmodbigint Safe);
+      ("%eqbigint", Pbigintcomp Ceq);
+      ("%noteqbigint", Pbigintcomp Cneq);
+      ("%ltbigint", Pbigintcomp Clt);
+      ("%lebigint", Pbigintcomp Cle);
+      ("%gtbigint", Pbigintcomp Cgt);
+      ("%gebigint", Pbigintcomp Cge);
       ("%string_length", Pstringlength);
       ("%string_safe_get", Pstringrefs);
       ("%string_unsafe_get", Pstringrefu);
@@ -402,7 +423,7 @@ let primitives_table =
 let find_primitive prim_name = Hashtbl.find primitives_table prim_name
 
 let specialize_comparison
-    ({ gencomp; intcomp; floatcomp; stringcomp; bytescomp; int64comp; boolcomp } :
+    ({ gencomp; intcomp; floatcomp; stringcomp; bytescomp; int64comp; bigintcomp; boolcomp } :
       specialized) env ty =
   match () with
   | ()
@@ -414,6 +435,7 @@ let specialize_comparison
   | () when is_base_type env ty Predef.path_string -> stringcomp
   | () when is_base_type env ty Predef.path_bytes -> bytescomp
   | () when is_base_type env ty Predef.path_int64 -> int64comp
+  | () when is_base_type env ty Predef.path_bigint -> bigintcomp
   | () when is_base_type env ty Predef.path_bool -> boolcomp
   | () -> gencomp
 
diff --git a/jscomp/runtime/caml.res b/jscomp/runtime/caml.res
@@ -52,6 +52,21 @@ let float_compare = (x: float, y: float) =>
     0
   }
 
+let bigint_compare = (x: bigint, y: bigint) =>
+  if x == y {
+    0
+  } else if x < y {
+    -1
+  } else if x > y {
+    1
+  } else if x == x {
+    1
+  } else if y == y {
+    -1
+  } else {
+    0
+  }
+
 /* Lexical order */
 let string_compare = (s1: string, s2: string): int =>
   if s1 == s2 {
diff --git a/jscomp/runtime/caml.resi b/jscomp/runtime/caml.resi
@@ -27,6 +27,7 @@ type selector<'a> = ('a, 'a) => 'a
 let int_compare: (int, int) => int
 let bool_compare: (bool, bool) => int
 let float_compare: (float, float) => int
+let bigint_compare: (bigint, bigint) => int
 let string_compare: (string, string) => int
 
 let bool_min: selector<bool>
diff --git a/lib/es6/caml.js b/lib/es6/caml.js
@@ -39,6 +39,20 @@ function float_compare(x, y) {
   }
 }
 
+function bigint_compare(x, y) {
+  if (x === y) {
+    return 0;
+  } else if (x < y) {
+    return -1;
+  } else if (x > y || x === x) {
+    return 1;
+  } else if (y === y) {
+    return -1;
+  } else {
+    return 0;
+  }
+}
+
 function string_compare(s1, s2) {
   if (s1 === s2) {
     return 0;
@@ -175,6 +189,7 @@ export {
   int_compare ,
   bool_compare ,
   float_compare ,
+  bigint_compare ,
   string_compare ,
   bool_min ,
   int_min ,
diff --git a/lib/js/caml.js b/lib/js/caml.js
@@ -39,6 +39,20 @@ function float_compare(x, y) {
   }
 }
 
+function bigint_compare(x, y) {
+  if (x === y) {
+    return 0;
+  } else if (x < y) {
+    return -1;
+  } else if (x > y || x === x) {
+    return 1;
+  } else if (y === y) {
+    return -1;
+  } else {
+    return 0;
+  }
+}
+
 function string_compare(s1, s2) {
   if (s1 === s2) {
     return 0;
@@ -174,6 +188,7 @@ function i64_max(x, y) {
 exports.int_compare = int_compare;
 exports.bool_compare = bool_compare;
 exports.float_compare = float_compare;
+exports.bigint_compare = bigint_compare;
 exports.string_compare = string_compare;
 exports.bool_min = bool_min;
 exports.int_min = int_min;
