Use the stdlib Map for the compiler's maps

CHANGELOG.md

@@ -61,6 +61,7 @@
 - Use the OCaml standard library's hash tables in the compiler. https://github.com/rescript-lang/rescript/pull/8786
 - Use the OCaml standard library's hash tables for the compiler's hash sets, and always order the imports of a module used both with and without `default` the same way. https://github.com/rescript-lang/rescript/pull/8787
 - Use the OCaml standard library's sets in the compiler. https://github.com/rescript-lang/rescript/pull/8789
+- Use the OCaml standard library's maps in the compiler. https://github.com/rescript-lang/rescript/pull/8790
 - Represent explicit expression braces as `Pexp_braces` in parsetree v1 and format `else` branches consistently with `if` branches. https://github.com/rescript-lang/rescript/pull/8678
 - Omit redundant braces around multi-statement switch case bodies when formatting. https://github.com/rescript-lang/rescript/pull/8677
 - Avoid running `rescript-schema-ppx` and `sury-ppx` on source files without an `@schema` annotation. https://github.com/rescript-lang/rescript/pull/8662

compiler/core/js_cmj_format.ml

@@ -67,12 +67,14 @@ let make ~(values : cmj_value Map_string.t) ~hoisted_exports ~effect_
     ~package_spec ~case : t =
   {
     values =
-      Map_string.to_sorted_array_with_f values (fun k v ->
+      Map_string.bindings values
+      |> List.map (fun (k, (v : cmj_value)) ->
           {
             name = k;
             arity = v.arity;
             persistent_closed_lambda = v.persistent_closed_lambda;
-          });
+          })
+      |> Array.of_list;
     hoisted_exports = Array.of_list hoisted_exports;
     pure = effect_ = None;
     package_spec;
@@ -108,7 +110,7 @@ let to_file name ~check_exists (v : t) =
     output_string oc s;
     close_out oc)
 
-let key_comp a b = Map_string.compare_key a b.name
+let key_comp a b = Ext_string.compare a b.name
 
 let not_found key =
   {name = key; arity = single_na; persistent_closed_lambda = None}

compiler/core/js_pass_tailcall_inline.ml

@@ -38,7 +38,14 @@ module S = Js_stmt_make
 let super = Js_record_map.super
 
 let substitue_variables (map : Ident.t Map_ident.t) =
-  {super with ident = (fun _ id -> Map_ident.find_default map id id)}
+  {
+    super with
+    ident =
+      (fun _ id ->
+        match Map_ident.find_opt id map with
+        | Some id -> id
+        | None -> id);
+  }
 
 (* 1. recursive value ? let rec x = 1 :: x
     non-terminating
@@ -64,13 +71,13 @@ let inline_call (immutable_list : bool list) params (args : J.expression list)
       Ext_list.fold_right2 params args (Map_ident.empty, processed_blocks)
         (fun param arg (map, acc) ->
           match arg.expression_desc with
-          | Var (Id id) -> (Map_ident.add map param id, acc)
+          | Var (Id id) -> (Map_ident.add param id map, acc)
           | _ -> (map, S.define_variable ~kind:Variable param arg :: acc))
     else
       Ext_list.fold_right3 params args immutable_list
         (Map_ident.empty, processed_blocks) (fun param arg mask (map, acc) ->
           match (mask, arg.expression_desc) with
-          | true, Var (Id id) -> (Map_ident.add map param id, acc)
+          | true, Var (Id id) -> (Map_ident.add param id map, acc)
           | _ -> (map, S.define_variable ~kind:Variable param arg :: acc))
   in
   if Map_ident.is_empty map then block

compiler/core/lam_beta_reduce.ml

@@ -97,7 +97,7 @@ let propagate_beta_reduce_with_map (meta : Lam_stats.t)
             ((p, arg) :: rest_bindings, Lambda.var p :: acc)
           | _ ->
             if Lam_analysis.no_side_effects arg then
-              match Map_ident.find_exn map old_param with
+              match Map_ident.find old_param map with
               | stat ->
                 if Lam_var_stats.top_and_used_zero_or_one stat then
                   (rest_bindings, arg :: acc)

compiler/core/lam_closure.ml

@@ -28,17 +28,19 @@ type stats = Lam_var_stats.stats
 
 let adjust (fv : stats Map_ident.t) (pos : position) (v : Ident.t) :
     stats Map_ident.t =
-  Map_ident.adjust fv v (fun v ->
+  Map_ident.update v
+    (fun v ->
       let stat =
         match v with
         | None -> Lam_var_stats.fresh_stats
         | Some v -> v
       in
-      Lam_var_stats.update stat pos)
+      Some (Lam_var_stats.update stat pos))
+    fv
 
 let param_map_of_list lst : stats Map_ident.t =
   Ext_list.fold_left lst Map_ident.empty (fun acc l ->
-      Map_ident.add acc l Lam_var_stats.fresh_stats)
+      Map_ident.add l Lam_var_stats.fresh_stats acc)
 
 (** Sanity check, remove all varaibles in [local_set] in the last pass *)
 let sink_pos = Lam_var_stats.sink
@@ -152,8 +154,9 @@ let free_variables (export_idents : Set_ident.t) (params : stats Map_ident.t)
 
 (** A bit consverative , it should be empty *)
 let is_closed lam =
-  Map_ident.for_all (free_variables Set_ident.empty Map_ident.empty lam)
+  Map_ident.for_all
     (fun k _ -> Ident.global k)
+    (free_variables Set_ident.empty Map_ident.empty lam)
 
 let is_closed_with_map (exports : Set_ident.t) (params : Ident.t list)
     (body : Lambda.t) : bool * stats Map_ident.t =

compiler/core/lam_coercion.ml

@@ -116,7 +116,7 @@ let handle_exports (meta : Lam_stats.t) (lambda_exports : Lambda.t list)
             {
               acc with
               export_list = newid :: acc.export_list;
-              export_map = Map_ident.add acc.export_map newid lam;
+              export_map = Map_ident.add newid lam acc.export_map;
               groups = Single (kind, newid, lam) :: acc.groups;
             }
         | _ ->
@@ -152,7 +152,7 @@ let handle_exports (meta : Lam_stats.t) (lambda_exports : Lambda.t list)
           {
             acc with
             export_list = newid :: acc.export_list;
-            export_map = Map_ident.add acc.export_map newid lam;
+            export_map = Map_ident.add newid lam acc.export_map;
             groups = Single (Strict, newid, lam) :: acc.groups;
           })
   in
@@ -162,7 +162,7 @@ let handle_exports (meta : Lam_stats.t) (lambda_exports : Lambda.t list)
       (fun (export_map, acc) x ->
         ( (match x with
           | Single (_, id, lam) when Set_ident.mem id export_set ->
-            Map_ident.add export_map id lam
+            Map_ident.add id lam export_map
           (* relies on the Invariant that [eoid] can not be bound before
               FIX: such invariant may not hold
           *)

compiler/core/lam_compile.ml

@@ -415,12 +415,15 @@ let compile output_prefix =
             ~return_unit ~async ~one_unit_arg ?directive
             ~immutable_mask:ret.immutable_mask
             (Ext_list.map params (fun x ->
-                 Map_ident.find_default ret.new_params x x))
+                 match Map_ident.find_opt x ret.new_params with
+                 | Some y -> y
+                 | None -> x))
             [
               S.while_ E.true_
-                (Map_ident.fold ret.new_params body_block
+                (Map_ident.fold
                    (fun old new_param acc ->
-                     S.define_variable ~kind:Alias old (E.var new_param) :: acc));
+                     S.define_variable ~kind:Alias old (E.var new_param) :: acc)
+                   ret.new_params body_block);
             ]
         else
           (* TODO:  save computation of length several times *)
@@ -1579,18 +1582,17 @@ let compile output_prefix =
                 (i + 1, assigns, new_params)
               | _ ->
                 let new_param, m =
-                  match Map_ident.find_opt ret.new_params param with
+                  match Map_ident.find_opt param ret.new_params with
                   | None ->
                     ret.immutable_mask.(i) <- false;
                     let v = Ext_ident.create ("_" ^ param.name) in
-                    (v, Map_ident.add new_params param v)
+                    (v, Map_ident.add param v new_params)
                   | Some v -> (v, new_params)
                 in
                 (i + 1, (new_param, arg) :: assigns, m))
         in
         ret.new_params <-
-          Map_ident.disjoint_merge_exn new_params ret.new_params (fun _ _ _ ->
-              assert false);
+          Map_ident.union (fun _ _ _ -> assert false) new_params ret.new_params;
         let block =
           Ext_list.map_append assigned_params
             [S.continue_ ()]

compiler/core/lam_compile_context.ml

@@ -117,16 +117,17 @@ let add_jmps (m : jmp_table) (exit_id : Ident.t) (code_table : handler list) :
     Ext_list.fold_left_with_offset code_table (m, [])
       (Handler_map.cardinal m + 1)
       (fun {label; handler; bindings} (acc, handlers) order_id ->
-        ( Handler_map.add acc label {exit_id; bindings; order_id},
+        ( Handler_map.add label {exit_id; bindings; order_id} acc,
           (order_id, handler) :: handlers ))
   in
   (map, List.rev handlers)
 
 let add_pseudo_jmp (m : jmp_table)
     (exit_id : Ident.t) (* TODO not needed, remove it later *)
     (code_table : handler) : jmp_table * Lambda.t =
-  ( Handler_map.add m code_table.label
-      {exit_id; bindings = code_table.bindings; order_id = -1},
+  ( Handler_map.add code_table.label
+      {exit_id; bindings = code_table.bindings; order_id = -1}
+      m,
     code_table.handler )
 
-let find_exn cxt i = Map_int.find_exn cxt.jmp_table i
+let find_exn cxt i = Map_int.find i cxt.jmp_table

compiler/core/lam_compile_main.ml

@@ -111,10 +111,10 @@ let js_hoisted_aliases (export_ids : Ident.t list)
     let group_map =
       Ext_list.fold_left groups Map_ident.empty (fun group_map group ->
           match group with
-          | Single (_, id, lam) -> Map_ident.add group_map id lam
+          | Single (_, id, lam) -> Map_ident.add id lam group_map
           | Recursive bindings ->
             Ext_list.fold_left bindings group_map (fun group_map (id, lam) ->
-                Map_ident.add group_map id lam)
+                Map_ident.add id lam group_map)
           | Nop _ -> group_map)
     in
     let rec access loc base fields =
@@ -131,7 +131,7 @@ let js_hoisted_aliases (export_ids : Ident.t list)
       | Lambda.Lvar id as lam -> (
         if Set_ident.mem id seen then (lam, Some id)
         else
-          match Map_ident.find_opt group_map id with
+          match Map_ident.find_opt id group_map with
           | Some
               ((Lambda.Lvar _ | Lambda.Lprim {primitive = Lambda.Pfield _; _})
                as alias) ->
@@ -172,7 +172,7 @@ let js_hoisted_aliases (export_ids : Ident.t list)
     in
     let exported_modules =
       Ext_list.fold_left export_ids Map_string.empty (fun modules id ->
-          Map_string.add modules id.Ident.name id)
+          Map_string.add id.Ident.name id modules)
     in
     let occupied_names =
       Ext_list.fold_left groups Set_string.empty (fun occupied group ->
@@ -195,9 +195,9 @@ let js_hoisted_aliases (export_ids : Ident.t list)
            in
            match path with
            | top :: fields -> (
-             match Map_string.find_opt exported_modules top with
+             match Map_string.find_opt top exported_modules with
              | Some top_id -> (
-               match Map_ident.find_opt group_map top_id with
+               match Map_ident.find_opt top_id group_map with
                | Some lam -> (
                  match find_path lam fields [] with
                  | Some (access_path, Some target_id, target)
@@ -352,7 +352,7 @@ let compile (output_prefix : string) export_idents hoisted (lam : Lambda.t) =
       ->
         ( group :: groups,
           id :: exports,
-          Map_ident.add export_map id lam,
+          Map_ident.add id lam export_map,
           {Js_cmj_format.path; export_name = name} :: hoisted_metadata ))
   in
   let groups = groups @ List.rev hoisted_groups in
@@ -366,8 +366,9 @@ let compile (output_prefix : string) export_idents hoisted (lam : Lambda.t) =
     }
   in
   let export_map =
-    Map_ident.fold hoisted_export_map coerced_input.export_map
-      (fun id lam acc -> Map_ident.add acc id lam)
+    Map_ident.fold
+      (fun id lam acc -> Map_ident.add id lam acc)
+      hoisted_export_map coerced_input.export_map
   in
   let () =
     if debug_ir then

compiler/core/lam_pass_count.ml

@@ -55,12 +55,12 @@ let collect_occurs lam : occ_tbl =
   let bind_var bv ident =
     let r = dummy_info () in
     Hash_ident.add occ ident r;
-    Map_ident.add bv ident r
+    Map_ident.add ident r bv
   in
 
   (* Record a use of a variable *)
   let add_one_use bv ident =
-    match Map_ident.find_opt bv ident with
+    match Map_ident.find_opt ident bv with
     | Some r -> r.times <- r.times + 1
     | None -> (
       (* ident is not locally bound, therefore this is a use under a lambda
@@ -79,7 +79,7 @@ let collect_occurs lam : occ_tbl =
       | None -> dummy_info ()
       | Some v -> v
     in
-    match Map_ident.find_opt bv ident with
+    match Map_ident.find_opt ident bv with
     | Some r -> absorb_info r n
     | None -> (
       (* ident is not locally bound, therefore this is a use under a lambda

compiler/core/lam_pass_deep_flatten.ml

@@ -30,7 +30,7 @@ let rec eliminate_tuple (id : Ident.t) (lam : Lambda.t) acc =
   match lam with
   | Llet (Alias, v, Lprim {primitive = Pfield (i, _); args = [Lvar tuple]}, e2)
     when Ident.same tuple id ->
-    eliminate_tuple id e2 (Map_int.add acc i v)
+    eliminate_tuple id e2 (Map_int.add i v acc)
   (* it is okay to have duplicates*)
   | _ -> if Lam_hit.hit_variable id lam then None else Some (acc, lam)
 (* [groups] are in reverse order *)
@@ -210,7 +210,7 @@ let deep_flatten (lam : Lambda.t) : Lambda.t =
         | Some (tuple_mapping, body) ->
           flatten
             (Ext_list.fold_left_with_offset args accux 0 (fun arg acc i ->
-                 match Map_int.find_opt tuple_mapping i with
+                 match Map_int.find_opt i tuple_mapping with
                  | None -> Lam_group.nop_cons arg acc
                  | Some key -> Lam_group.single str key arg :: acc))
             body

compiler/core/lam_stats_export.ml

@@ -41,7 +41,7 @@ let values_of_export (meta : Lam_stats.t) (export_map : Lambda.t Map_ident.t) :
                 | NA -> Lam_arity.na
                 | SimpleForm lam -> Lam_arity_analysis.get_arity meta lam))
         | Some _ | None -> (
-          match Map_ident.find_opt export_map x with
+          match Map_ident.find_opt x export_map with
           | Some (Lprim {primitive = Pmakeblock info; args})
             when Lambda.is_immutable_block info ->
             Submodule
@@ -50,7 +50,7 @@ let values_of_export (meta : Lam_stats.t) (export_map : Lambda.t Map_ident.t) :
           | Some _ | None -> single_na)
       in
       let persistent_closed_lambda =
-        let optlam = Map_ident.find_opt export_map x in
+        let optlam = Map_ident.find_opt x export_map in
         match optlam with
         | Some
             (Lconst
@@ -105,7 +105,7 @@ let values_of_export (meta : Lam_stats.t) (export_map : Lambda.t Map_ident.t) :
         let cmj_value : Js_cmj_format.cmj_value =
           {arity; persistent_closed_lambda}
         in
-        Map_string.add acc x.name cmj_value)
+        Map_string.add x.name cmj_value acc)
 
 (* ATTENTION: all runtime modules, if it is not hard required,
    it should be okay to not reference it