Unified operators

CHANGELOG.md

@@ -12,6 +12,10 @@
 
 # 12.0.0-alpha.5 (Unreleased)
 
+#### :rocket: New Feature
+
+- Introduce "Unified operators" for arithmetic operators (`+`, `-`, `*`, `/`, `mod`). See https://github.com/rescript-lang/rescript-compiler/pull/7057
+
 # 12.0.0-alpha.4
 
 #### :boom: Breaking Change

compiler/ml/translcore.ml

@@ -49,6 +49,33 @@ let transl_extension_constructor env path ext =
 
 (* Translation of primitives *)
 
+(** This is ad-hoc translation for unifying specific primitive operations
+     See [Unified_ops] module for detailed explanation.
+  *)
+let translate_unified_ops (prim : Primitive.description) (env : Env.t)
+    (lhs_type : type_expr) : Lambda.primitive option =
+  (* lhs_type is already unified in type-level *)
+  let entry = Hashtbl.find_opt Unified_ops.index_by_name prim.prim_name in
+  match entry with
+  | Some {specialization} -> (
+    match specialization with
+    | {int}
+      when is_base_type env lhs_type Predef.path_int
+           || maybe_pointer_type env lhs_type = Immediate ->
+      Some int
+    | {float = Some float} when is_base_type env lhs_type Predef.path_float ->
+      Some float
+    | {bigint = Some bigint} when is_base_type env lhs_type Predef.path_bigint
+      ->
+      Some bigint
+    | {string = Some string} when is_base_type env lhs_type Predef.path_string
+      ->
+      Some string
+    | {bool = Some bool} when is_base_type env lhs_type Predef.path_bool ->
+      Some bool
+    | {int} -> Some int)
+  | _ -> None
+
 type specialized = {
   objcomp: Lambda.primitive;
   intcomp: Lambda.primitive;
@@ -394,12 +421,21 @@ let specialize_comparison
    raise Not_found if primitive is unknown *)
 
 let specialize_primitive p env ty (* ~has_constant_constructor *) =
-  try
-    let table = Hashtbl.find comparisons_table p.prim_name in
-    match is_function_type env ty with
-    | Some (lhs, _rhs) -> specialize_comparison table env lhs
-    | None -> table.objcomp
-  with Not_found -> find_primitive p.prim_name
+  let fn_expr = is_function_type env ty in
+  let unified =
+    match fn_expr with
+    | Some (lhs, _) -> translate_unified_ops p env lhs
+    | None -> None
+  in
+  match unified with
+  | Some primitive -> primitive
+  | None -> (
+    try
+      let table = Hashtbl.find comparisons_table p.prim_name in
+      match fn_expr with
+      | Some (lhs, _rhs) -> specialize_comparison table env lhs
+      | None -> table.objcomp
+    with Not_found -> find_primitive p.prim_name)
 
 (* Eta-expand a primitive *)
 
@@ -458,32 +494,44 @@ let transl_primitive loc p env ty =
 
 let transl_primitive_application loc prim env ty args =
   let prim_name = prim.prim_name in
-  try
+  let unified =
     match args with
-    | [arg1; _]
-      when is_base_type env arg1.exp_type Predef.path_bool
-           && Hashtbl.mem comparisons_table prim_name ->
-      (Hashtbl.find comparisons_table prim_name).boolcomp
-    | _ ->
-      let has_constant_constructor =
-        match args with
-        | [_; {exp_desc = Texp_construct (_, {cstr_tag = Cstr_constant _}, _)}]
-        | [{exp_desc = Texp_construct (_, {cstr_tag = Cstr_constant _}, _)}; _]
-        | [_; {exp_desc = Texp_variant (_, None)}]
-        | [{exp_desc = Texp_variant (_, None)}; _] ->
-          true
-        | _ -> false
-      in
-      if has_constant_constructor then
-        match Hashtbl.find_opt comparisons_table prim_name with
-        | Some table when table.simplify_constant_constructor -> table.intcomp
-        | Some _ | None -> specialize_primitive prim env ty
-        (* ~has_constant_constructor*)
-      else specialize_primitive prim env ty
-  with Not_found ->
-    if String.length prim_name > 0 && prim_name.[0] = '%' then
-      raise (Error (loc, Unknown_builtin_primitive prim_name));
-    Pccall prim
+    | [arg1] | [arg1; _] -> translate_unified_ops prim env arg1.exp_type
+    | _ -> None
+  in
+  match unified with
+  | Some primitive -> primitive
+  | None -> (
+    try
+      match args with
+      | [arg1; _]
+        when is_base_type env arg1.exp_type Predef.path_bool
+             && Hashtbl.mem comparisons_table prim_name ->
+        (Hashtbl.find comparisons_table prim_name).boolcomp
+      | _ ->
+        let has_constant_constructor =
+          match args with
+          | [
+              _; {exp_desc = Texp_construct (_, {cstr_tag = Cstr_constant _}, _)};
+            ]
+          | [
+              {exp_desc = Texp_construct (_, {cstr_tag = Cstr_constant _}, _)}; _;
+            ]
+          | [_; {exp_desc = Texp_variant (_, None)}]
+          | [{exp_desc = Texp_variant (_, None)}; _] ->
+            true
+          | _ -> false
+        in
+        if has_constant_constructor then
+          match Hashtbl.find_opt comparisons_table prim_name with
+          | Some table when table.simplify_constant_constructor -> table.intcomp
+          | Some _ | None -> specialize_primitive prim env ty
+          (* ~has_constant_constructor*)
+        else specialize_primitive prim env ty
+    with Not_found ->
+      if String.length prim_name > 0 && prim_name.[0] = '%' then
+        raise (Error (loc, Unknown_builtin_primitive prim_name));
+      Pccall prim)
 
 (* To propagate structured constants *)
 

compiler/ml/typecore.ml

@@ -2458,7 +2458,9 @@ and type_expect_ ?type_clash_context ?in_function ?(recarg = Rejected) env sexp
     in
     let type_clash_context = type_clash_context_from_function sexp sfunct in
     let args, ty_res, fully_applied =
-      type_application ?type_clash_context uncurried env funct sargs
+      match translate_unified_ops env funct sargs with
+      | Some (targs, result_type) -> (targs, result_type, true)
+      | None -> type_application ?type_clash_context uncurried env funct sargs
     in
     end_def ();
     unify_var env (newvar ()) funct.exp_type;
@@ -3561,6 +3563,101 @@ and is_automatic_curried_application env funct =
   | Tarrow _ -> true
   | _ -> false
 
+(** This is ad-hoc translation for unifying specific primitive operations
+     See [Unified_ops] module for detailed explanation.
+  *)
+and translate_unified_ops (env : Env.t) (funct : Typedtree.expression)
+    (sargs : sargs) : (targs * Types.type_expr) option =
+  match funct.exp_desc with
+  | Texp_ident (path, _, _) -> (
+    let entry = Hashtbl.find_opt Unified_ops.index_by_path (Path.name path) in
+    match (entry, sargs) with
+    | Some {form = Unary; specialization}, [(lhs_label, lhs_expr)] ->
+      let lhs = type_exp env lhs_expr in
+      let lhs_type = expand_head env lhs.exp_type in
+      let result_type =
+        match (lhs_type.desc, specialization) with
+        | Tconstr (path, _, _), _ when Path.same path Predef.path_int ->
+          Predef.type_int
+        | Tconstr (path, _, _), {bool = Some _}
+          when Path.same path Predef.path_bool ->
+          Predef.type_bool
+        | Tconstr (path, _, _), {float = Some _}
+          when Path.same path Predef.path_float ->
+          Predef.type_float
+        | Tconstr (path, _, _), {bigint = Some _}
+          when Path.same path Predef.path_bigint ->
+          Predef.type_bigint
+        | Tconstr (path, _, _), {string = Some _}
+          when Path.same path Predef.path_string ->
+          Predef.type_string
+        | _ ->
+          unify env lhs_type Predef.type_int;
+          Predef.type_int
+      in
+      let targs = [(lhs_label, Some lhs)] in
+      Some (targs, result_type)
+    | ( Some {form = Binary; specialization},
+        [(lhs_label, lhs_expr); (rhs_label, rhs_expr)] ) ->
+      let lhs = type_exp env lhs_expr in
+      let lhs_type = expand_head env lhs.exp_type in
+      let rhs = type_exp env rhs_expr in
+      let rhs_type = expand_head env rhs.exp_type in
+      let lhs, rhs, result_type =
+        (* Rule 1. Try unifying to lhs *)
+        match (lhs_type.desc, specialization) with
+        | Tconstr (path, _, _), _ when Path.same path Predef.path_int ->
+          let rhs = type_expect env rhs_expr Predef.type_int in
+          (lhs, rhs, Predef.type_int)
+        | Tconstr (path, _, _), {bool = Some _}
+          when Path.same path Predef.path_bool ->
+          let rhs = type_expect env rhs_expr Predef.type_bool in
+          (lhs, rhs, Predef.type_bool)
+        | Tconstr (path, _, _), {float = Some _}
+          when Path.same path Predef.path_float ->
+          let rhs = type_expect env rhs_expr Predef.type_float in
+          (lhs, rhs, Predef.type_float)
+        | Tconstr (path, _, _), {bigint = Some _}
+          when Path.same path Predef.path_bigint ->
+          let rhs = type_expect env rhs_expr Predef.type_bigint in
+          (lhs, rhs, Predef.type_bigint)
+        | Tconstr (path, _, _), {string = Some _}
+          when Path.same path Predef.path_string ->
+          let rhs = type_expect env rhs_expr Predef.type_string in
+          (lhs, rhs, Predef.type_string)
+        | _ -> (
+          (* Rule 2. Try unifying to rhs *)
+          match (rhs_type.desc, specialization) with
+          | Tconstr (path, _, _), _ when Path.same path Predef.path_int ->
+            let lhs = type_expect env lhs_expr Predef.type_int in
+            (lhs, rhs, Predef.type_int)
+          | Tconstr (path, _, _), {bool = Some _}
+            when Path.same path Predef.path_bool ->
+            let lhs = type_expect env lhs_expr Predef.type_bool in
+            (lhs, rhs, Predef.type_bool)
+          | Tconstr (path, _, _), {float = Some _}
+            when Path.same path Predef.path_float ->
+            let lhs = type_expect env lhs_expr Predef.type_float in
+            (lhs, rhs, Predef.type_float)
+          | Tconstr (path, _, _), {bigint = Some _}
+            when Path.same path Predef.path_bigint ->
+            let lhs = type_expect env lhs_expr Predef.type_bigint in
+            (lhs, rhs, Predef.type_bigint)
+          | Tconstr (path, _, _), {string = Some _}
+            when Path.same path Predef.path_string ->
+            let lhs = type_expect env lhs_expr Predef.type_string in
+            (lhs, rhs, Predef.type_string)
+          | _ ->
+            (* Rule 3. Fallback to int *)
+            let lhs = type_expect env lhs_expr Predef.type_int in
+            let rhs = type_expect env rhs_expr Predef.type_int in
+            (lhs, rhs, Predef.type_int))
+      in
+      let targs = [(lhs_label, Some lhs); (rhs_label, Some rhs)] in
+      Some (targs, result_type)
+    | _ -> None)
+  | _ -> None
+
 and type_application ?type_clash_context uncurried env funct (sargs : sargs) :
     targs * Types.type_expr * bool =
   (* funct.exp_type may be generic *)