Merge pull request #3586 from FStarLang/_nik_cache_dsenv_lookup

Optimize name resolution in desugaring

src/syntax/FStarC.Syntax.DsEnv.fst

@@ -45,11 +45,11 @@ type rec_binding = (ident & lid & (* name bound by recur
  used_marker) (* this ref marks whether it was used, so we can warn if not *)
 
 type scope_mod =
-| Local_binding  of local_binding
+| Local_bindings of BU.psmap local_binding (* a map local bindings in a scope; a map to avoid a linear scan *) 
 | Rec_binding of rec_binding
 | Module_abbrev of module_abbrev
 | Open_module_or_namespace of open_module_or_namespace
-| Top_level_def  of ident   (* top-level definition for an unqualified identifier x to be resolved as curmodule.x. *)
+| Top_level_defs of BU.psmap bool (* a map (to avoid a linear scan) recording that a top-level definition for an unqualified identifier x is in scope and should be resolved as curmodule.x. *)
 | Record_or_dc of record_or_dc (* to honor interleavings of "open" and record definitions *)
 
 type string_set = RBSet.t string
@@ -303,8 +303,9 @@ let try_lookup_id''
  in
  let curmod_ns = ids_of_lid (current_module env) in
  let proc = function
- | Local_binding l
- when check_local_binding_id l ->
+ | Local_bindings lbs
+ when Some? (BU.psmap_try_find lbs (string_of_id id)) ->
+ let Some l = BU.psmap_try_find lbs (string_of_id id) in
  let (_, _, used_marker) = l in
  used_marker := true;
  k_local_binding l
@@ -320,8 +321,8 @@ let try_lookup_id''
  | None -> Cont_ignore
  | Some id -> find_in_module_with_includes eikind find_in_module Cont_ignore env ns id)
 
- | Top_level_def id'
- when string_of_id id' = string_of_id id ->
+ | Top_level_defs ids
+ when Some? (BU.psmap_try_find ids (string_of_id id)) ->
  (* indicates a global definition shadowing previous
  "open"s. If the definition is not actually found by the
  [lookup_default_id] finder, then it may mean that we are in a
@@ -982,7 +983,14 @@ let push_bv' env (x:ident) =
  let r = range_of_id x in
  let bv = S.gen_bv (string_of_id x) (Some r) ({ tun with pos = r }) in
  let used_marker = BU.mk_ref false in
- push_scope_mod env (Local_binding (x, bv, used_marker)), bv, used_marker
+ let scope_mods =
+ match env.scope_mods with
+ | Local_bindings lbs :: rest ->
+ Local_bindings (BU.psmap_add lbs (string_of_id x) (x, bv, used_marker)) :: rest
+ | _ ->
+ Local_bindings (BU.psmap_add (BU.psmap_empty()) (string_of_id x) (x, bv, used_marker)) :: env.scope_mods
+ in
+ { env with scope_mods }, bv, used_marker
 
 let push_bv env x =
  let (env, bv, _) = push_bv' env x in
@@ -1027,12 +1035,19 @@ let push_sigelt' fail_on_dup env s =
  let env, lss = match s.sigel with
  | Sig_bundle {ses} -> env, List.map (fun se -> (lids_of_sigelt se, se)) ses
  | _ -> env, [lids_of_sigelt s, s] in
+ let push_top_level_def id stack =
+ match stack with
+ | Top_level_defs ids :: rest ->
+ Top_level_defs (BU.psmap_add ids (string_of_id id) true) :: rest
+ | _ ->
+ Top_level_defs (BU.psmap_add (BU.psmap_empty()) (string_of_id id) true) :: stack
+ in
  lss |> List.iter (fun (lids, se) ->
  lids |> List.iter (fun lid ->
  (* the identifier is added into the list of global
  declarations, to allow shadowing of definitions that were
  formerly reachable by previous "open"s. *)
- let () = globals := Top_level_def (ident_of_lid lid) :: !globals in
+ let () = globals := push_top_level_def (ident_of_lid lid) !globals in
  (* the identifier is added into the list of global identifiers
  of the corresponding module to shadow any "include" *)
  let modul = string_of_lid (lid_of_ids (ns_of_lid lid)) in