From 9ef2345aec8de5415a51d97a2e56fb0f2dcd7c32 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 4 Aug 2021 18:53:13 +0200
Subject: [PATCH] perf: enforce hash map load factor

---
 src/Lean/Environment.lean |  5 ++---
 src/Std/Data/HashMap.lean | 15 ++++++++++-----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean
index 4736872daf..bd5fbee994 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -614,10 +614,9 @@ partial def importModules (imports : List Import) (opts : Options) (trustLevel :
     let mut numConsts := 0
     for mod in s.moduleData do
       numConsts := numConsts + mod.constants.size
-    -- (moduleNames, mods, regions)
     let mut modIdx : Nat := 0
-    let mut const2ModIdx : HashMap Name ModuleIdx := Std.mkHashMap (nbuckets := numConsts)
-    let mut constantMap : HashMap Name ConstantInfo := Std.mkHashMap (nbuckets := numConsts)
+    let mut const2ModIdx : HashMap Name ModuleIdx := Std.mkHashMap (capacity := numConsts)
+    let mut constantMap : HashMap Name ConstantInfo := Std.mkHashMap (capacity := numConsts)
     for mod in s.moduleData do
       for cinfo in mod.constants do
         const2ModIdx := const2ModIdx.insert cinfo.name modIdx
diff --git a/src/Std/Data/HashMap.lean b/src/Std/Data/HashMap.lean
index 7cf5b079a1..22ab2b4d30 100644
--- a/src/Std/Data/HashMap.lean
+++ b/src/Std/Data/HashMap.lean
@@ -18,8 +18,13 @@ structure HashMapImp (α : Type u) (β : Type v) where
   size       : Nat
   buckets    : HashMapBucket α β
 
-def mkHashMapImp {α : Type u} {β : Type v} (nbuckets := 8) : HashMapImp α β :=
-  let n := if nbuckets = 0 then 8 else nbuckets;
+private def numBucketsForCapacity (capacity : Nat) : Nat :=
+  -- a "load factor" of 0.75 is the usual standard for hash maps
+  capacity * 4 / 3
+
+def mkHashMapImp {α : Type u} {β : Type v} (capacity := 0) : HashMapImp α β :=
+  let_fun nbuckets := numBucketsForCapacity capacity
+  let n := if nbuckets = 0 then 8 else nbuckets
   { size       := 0,
     buckets    :=
     ⟨ mkArray n AssocList.nil,
@@ -99,7 +104,7 @@ def expand [Hashable α] (size : Nat) (buckets : HashMapBucket α β) : HashMapI
     else
       let size'    := size + 1
       let buckets' := buckets.update i (AssocList.cons a b bkt) h
-      if size' ≤ buckets.val.size then
+      if numBucketsForCapacity size' ≤ buckets.val.size then
         ({ size := size', buckets := buckets' }, false)
       else
         (expand size' buckets', false)
@@ -124,8 +129,8 @@ def HashMap (α : Type u) (β : Type v) [BEq α] [Hashable α] :=
 
 open Std.HashMapImp
 
-def mkHashMap {α : Type u} {β : Type v} [BEq α] [Hashable α] (nbuckets := 8) : HashMap α β :=
-  ⟨ mkHashMapImp nbuckets, WellFormed.mkWff nbuckets ⟩
+def mkHashMap {α : Type u} {β : Type v} [BEq α] [Hashable α] (capacity := 0) : HashMap α β :=
+  ⟨ mkHashMapImp capacity, WellFormed.mkWff capacity ⟩
 
 namespace HashMap
 instance [BEq α] [Hashable α] : Inhabited (HashMap α β) where