feat(library/init/lean/environment): add Environment.displayStats and --stats command line argument

library/init/data/rbmap/basic.lean

@@ -293,6 +293,9 @@ l.foldl (λ r p, r.insert p.1 p.2) (mkRBMap α β lt)
 @[inline] def any : RBMap α β lt → (α → β → Bool) → Bool
 | ⟨t, _⟩ p := t.any p
 
+def size (m : RBMap α β lt) : Nat :=
+m.fold (λ sz _ _, sz+1) 0
+
 end RBMap
 
 def rbmapOf {α : Type u} {β : Type v} (l : List (α × β)) (lt : α → α → Bool) : RBMap α β lt :=

library/init/lean/environment.lean

@@ -408,4 +408,30 @@ env ← finalizePersistentExtensions env,
 env ← mods.miterate env $ λ _ mod env, performModifications env mod.serialized,
 pure env
 
+namespace Environment
+
+@[export lean.display_stats_core]
+def displayStats (env : Environment) : IO Unit :=
+do
+pExtDescrs ← persistentEnvExtensionsRef.get,
+let numModules := ((pExtDescrs.get 0).toEnvExtension.getState env).importedEntries.size,
+IO.println ("direct imports:             " ++ toString env.imports),
+IO.println ("number of imported modules: " ++ toString numModules),
+IO.println ("number of consts:           " ++ toString env.constants.size),
+IO.println ("number of imported consts:  " ++ toString env.constants.stageSizes.1),
+IO.println ("number of local consts:     " ++ toString env.constants.stageSizes.2),
+IO.println ("trust level:                " ++ toString env.trustLevel),
+IO.println ("number of extensions:       " ++ toString env.extensions.size),
+pExtDescrs.mfor $ λ extDescr, do {
+  IO.println ("extension '" ++ toString extDescr.name ++ "'"),
+  let s := extDescr.toEnvExtension.getState env,
+  IO.println ("  lazy:                       " ++ toString extDescr.lazy),
+  IO.println ("  number of imported entries: " ++ toString (s.importedEntries.foldl (λ sum es, sum + es.size) 0)),
+  IO.println ("  number of local entries:    " ++ toString s.entries.length),
+  IO.println ("  forced state:               " ++ toString s.state.isSome),
+  pure ()
+},
+pure ()
+
+end Environment
 end Lean

library/init/lean/smap.lean

@@ -55,5 +55,11 @@ if m.stage₁ then { stage₁ := false, .. m } else m
 @[inline] def foldStage2 {σ : Type w} (f : σ → α → β → σ) (s : σ) (m : SMap α β lt) : σ :=
 m.map₂.fold f s
 
+def size (m : SMap α β lt) : Nat :=
+m.map₁.size + m.map₂.size
+
+def stageSizes (m : SMap α β lt) : Nat × Nat :=
+(m.map₁.size, m.map₂.size)
+
 end SMap
 end Lean

src/kernel/environment.cpp

@@ -297,6 +297,12 @@ void environment::for_each_constant(std::function<void(constant_info const & d)>
         });
 }
 
+obj_res display_stats_core(obj_arg env, obj_arg w);
+
+void environment::display_stats() const {
+    dec_ref(display_stats_core(get_obj_arg(), io_mk_world()));
+}
+
 void initialize_environment() {
     g_env_ext_class = register_external_object_class(env_ext_finalizer, env_ext_foreach);
 }

src/kernel/environment.h

@@ -94,6 +94,8 @@ public:
     friend bool is_eqp(environment const & e1, environment const & e2) {
         return e1.raw() == e2.raw();
     }
+
+    void display_stats() const;
 };
 
 void check_no_metavar_no_fvar(environment const & env, name const & n, expr const & e);

src/shell/lean.cpp

@@ -203,6 +203,7 @@ static void display_help(std::ostream & out) {
 #endif
     std::cout << "  --new-frontend     use new frontend\n";
     std::cout << "  --profile          display elaboration/type checking time for each definition/theorem\n";
+    std::cout << "  --stats            display environment statistics\n";
     DEBUG_CODE(
     std::cout << "  --debug=tag        enable assertions with the given tag\n";
         )
@@ -218,6 +219,7 @@ static struct option g_long_options[] = {
     {"memory",       required_argument, 0, 'M'},
     {"trust",        required_argument, 0, 't'},
     {"profile",      no_argument,       0, 'P'},
+    {"stats",        no_argument,       0, 'a'},
     {"threads",      required_argument, 0, 'j'},
     {"quiet",        no_argument,       0, 'q'},
     {"deps",         no_argument,       0, 'd'},
@@ -239,7 +241,7 @@ static struct option g_long_options[] = {
 };
 
 static char const * g_opt_str =
-    "PdD:c:qpgvhet:012E:A:B:j:012rM:012T:012"
+    "Pdc:qpgvht:012j:012rM:012T:012a"
 #if defined(LEAN_MULTI_THREAD)
     "s:012"
 #endif
@@ -317,6 +319,7 @@ int main(int argc, char ** argv) {
     unsigned trust_lvl = LEAN_BELIEVER_TRUST_LEVEL + 1;
     bool only_deps = false;
     bool new_frontend = false;
+    bool stats = false;
     // unsigned num_threads    = 0;
 #if defined(LEAN_MULTI_THREAD)
     // num_threads = hardware_concurrency();
@@ -376,6 +379,9 @@ int main(int argc, char ** argv) {
             case 'd':
                 only_deps = true;
                 break;
+            case 'a':
+                stats = true;
+                break;
             case 'D':
                 try {
                     opts = set_config_option(opts, optarg);
@@ -538,6 +544,10 @@ int main(int argc, char ** argv) {
             env = p.env();
         }
 
+        if (stats) {
+            env.display_stats();
+        }
+
         if (make_mode && ok) {
             auto olean_fn = olean_of_lean(mod_fn);
             time_task t(".olean serialization",

src/stage0/init/data/rbmap/basic.cpp

@@ -42,6 +42,7 @@ obj* l_RBNode_insert___boxed(obj*, obj*);
 obj* l_RBMap_any___main(obj*, obj*, obj*);
 obj* l_RBMap_contains___boxed(obj*, obj*);
 obj* l_RBMap_any___main___rarg___boxed(obj*, obj*);
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___rarg___boxed(obj*, obj*);
 obj* l_RBNode_setBlack(obj*, obj*);
 obj* l_List_foldl___main___at_rbmapOf___spec__1(obj*, obj*);
 obj* l_RBNode_isRed___boxed(obj*, obj*);
@@ -88,6 +89,7 @@ obj* l_RBMap_toList___boxed(obj*, obj*, obj*);
 obj* l_RBNode_min___main___rarg(obj*);
 obj* l_RBMap_mfold(obj*, obj*, obj*, obj*, obj*);
 obj* l_RBNode_ins___boxed(obj*, obj*);
+obj* l_RBMap_size(obj*, obj*, obj*);
 obj* l_RBMap_mfor(obj*, obj*, obj*, obj*, obj*);
 obj* l_RBNode_ins___main___rarg(obj*, obj*, obj*, obj*);
 obj* l_RBNode_lowerBound___main(obj*, obj*);
@@ -102,10 +104,12 @@ obj* l_RBNode_max(obj*, obj*);
 obj* l_RBNode_mfold___main___at_RBMap_mfor___spec__1(obj*, obj*, obj*, obj*);
 obj* l_RBNode_isBlack___rarg___boxed(obj*);
 obj* l_RBMap_empty___boxed(obj*, obj*, obj*);
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1(obj*, obj*);
 obj* l_RBNode_mfold___main___rarg(obj*, obj*, obj*, obj*);
 obj* l_RBNode_isRed(obj*, obj*);
 obj* l_RBMap_toList___main(obj*, obj*, obj*);
 uint8 l_RBMap_any___main___rarg(obj*, obj*);
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(obj*, obj*);
 obj* l_RBNode_fold___main___boxed(obj*, obj*, obj*);
 obj* l_RBNode_min___rarg(obj*);
 obj* l_RBMap_findCore___main___rarg(obj*, obj*, obj*);
@@ -265,9 +269,11 @@ obj* l_RBNode_balance1___boxed(obj*, obj*);
 obj* l_RBNode_depth___main___boxed(obj*, obj*);
 obj* l_RBNode_isBlack(obj*, obj*);
 obj* l_RBMap_revFold___main(obj*, obj*, obj*, obj*);
+obj* l_RBMap_size___rarg___boxed(obj*);
 obj* l_RBMap_all___main___rarg___boxed(obj*, obj*);
 obj* l_RBNode_find___rarg___boxed(obj*, obj*, obj*, obj*);
 obj* l_RBNode_balance_u_2083___main___boxed(obj*, obj*);
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___boxed(obj*, obj*);
 obj* l_RBMap_fromList(obj*, obj*);
 obj* l_RBNode_any___rarg___boxed(obj*, obj*);
 obj* l_RBNode_max___main(obj*, obj*);
@@ -293,6 +299,7 @@ obj* l_RBNode_ins___rarg(obj*, obj*, obj*, obj*);
 obj* l_RBMap_all___rarg___boxed(obj*, obj*);
 obj* l_RBNode_setRed___rarg(obj*);
 obj* l_RBNode_fold___boxed(obj*, obj*, obj*);
+obj* l_RBMap_size___rarg(obj*);
 obj* l_RBNode_fold___main(obj*, obj*, obj*);
 obj* l_List_foldl___main___at_rbmapOf___spec__1___boxed(obj*, obj*);
 obj* l_RBNode_fold(obj*, obj*, obj*);
@@ -309,6 +316,7 @@ obj* l_RBNode_setRed___boxed(obj*, obj*);
 obj* l_RBNode_find___main___rarg___boxed(obj*, obj*, obj*, obj*);
 obj* l_RBMap_insert___main___rarg(obj*, obj*, obj*, obj*);
 uint8 l_RBMap_any___rarg(obj*, obj*);
+obj* l_RBMap_size___boxed(obj*, obj*, obj*);
 obj* l_RBMap_any(obj*, obj*, obj*);
 obj* l_RBNode_setRed(obj*, obj*);
 obj* l_RBNode_balance_u_2083(obj*, obj*);
@@ -13542,6 +13550,92 @@ lean::dec(x_2);
 return x_3;
 }
 }
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(obj* x_0, obj* x_1) {
+_start:
+{
+if (lean::obj_tag(x_1) == 0)
+{
+return x_0;
+}
+else
+{
+obj* x_2; obj* x_3; obj* x_4; obj* x_5; obj* x_6; 
+x_2 = lean::cnstr_get(x_1, 0);
+x_3 = lean::cnstr_get(x_1, 3);
+x_4 = l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(x_0, x_2);
+x_5 = lean::mk_nat_obj(1ul);
+x_6 = lean::nat_add(x_4, x_5);
+lean::dec(x_4);
+x_0 = x_6;
+x_1 = x_3;
+goto _start;
+}
+}
+}
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1(obj* x_0, obj* x_1) {
+_start:
+{
+obj* x_2; 
+x_2 = lean::alloc_closure(reinterpret_cast<void*>(l_RBNode_fold___main___at_RBMap_size___spec__1___rarg___boxed), 2, 0);
+return x_2;
+}
+}
+obj* l_RBMap_size___rarg(obj* x_0) {
+_start:
+{
+obj* x_1; obj* x_2; 
+x_1 = lean::mk_nat_obj(0ul);
+x_2 = l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(x_1, x_0);
+return x_2;
+}
+}
+obj* l_RBMap_size(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = lean::alloc_closure(reinterpret_cast<void*>(l_RBMap_size___rarg___boxed), 1, 0);
+return x_3;
+}
+}
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___rarg___boxed(obj* x_0, obj* x_1) {
+_start:
+{
+obj* x_2; 
+x_2 = l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(x_0, x_1);
+lean::dec(x_1);
+return x_2;
+}
+}
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___boxed(obj* x_0, obj* x_1) {
+_start:
+{
+obj* x_2; 
+x_2 = l_RBNode_fold___main___at_RBMap_size___spec__1(x_0, x_1);
+lean::dec(x_0);
+lean::dec(x_1);
+return x_2;
+}
+}
+obj* l_RBMap_size___rarg___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_RBMap_size___rarg(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
+obj* l_RBMap_size___boxed(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = l_RBMap_size(x_0, x_1, x_2);
+lean::dec(x_0);
+lean::dec(x_1);
+lean::dec(x_2);
+return x_3;
+}
+}
 obj* l_List_foldl___main___at_rbmapOf___spec__1___rarg(obj* x_0, obj* x_1, obj* x_2) {
 _start:
 {

src/stage0/init/lean/smap.cpp

@@ -25,6 +25,8 @@ obj* l_Lean_SMap_contains(obj*, obj*);
 obj* l_Lean_SMap_insert___boxed(obj*, obj*);
 obj* l_Lean_SMap_HasEmptyc___boxed(obj*, obj*);
 obj* l_Lean_SMap_find___rarg(obj*, obj*, obj*, obj*, obj*);
+obj* l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(obj*, obj*);
+obj* l_Lean_SMap_stageSizes(obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_foldStage2___boxed(obj*, obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_empty___boxed(obj*, obj*);
 obj* l_Lean_SMap_find___main___rarg(obj*, obj*, obj*, obj*, obj*);
@@ -33,14 +35,20 @@ obj* l_Lean_SMap_insert___main___rarg(obj*, obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_insert(obj*, obj*);
 obj* l_Lean_SMap_insert___main(obj*, obj*);
 obj* l_Lean_SMap_HasEmptyc(obj*, obj*);
+obj* l_Lean_SMap_size(obj*, obj*, obj*, obj*, obj*);
+obj* l_Lean_SMap_size___rarg___boxed(obj*);
 obj* l_Lean_SMap_contains___main(obj*, obj*);
 obj* l_Lean_SMap_contains___boxed(obj*, obj*);
 uint8 l_Lean_SMap_contains___main___rarg(obj*, obj*, obj*, obj*, obj*);
+namespace lean {
+obj* nat_add(obj*, obj*);
+}
 obj* l_HashMapImp_find___rarg(obj*, obj*, obj*, obj*);
 obj* l_HashMapImp_insert___rarg(obj*, obj*, obj*, obj*, obj*);
 obj* l_RBNode_fold___main___rarg(obj*, obj*, obj*);
 obj* l_Lean_SMap_switch(obj*, obj*);
 obj* l_Lean_SMap_contains___rarg___boxed(obj*, obj*, obj*, obj*, obj*);
+obj* l_Lean_SMap_size___rarg(obj*);
 obj* l_RBNode_find___main___rarg(obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_empty(obj*, obj*);
 obj* l_Lean_SMap_find___main___boxed(obj*, obj*);
@@ -48,14 +56,17 @@ obj* l_Lean_SMap_insert___main___boxed(obj*, obj*);
 uint8 l_Lean_SMap_contains___rarg(obj*, obj*, obj*, obj*, obj*);
 extern obj* l_HashMap_Inhabited___closed__1;
 obj* l_Lean_SMap_find___main(obj*, obj*);
+obj* l_Lean_SMap_stageSizes___boxed(obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_HasEmptyc___rarg___boxed(obj*, obj*, obj*);
 obj* l_Lean_SMap_HasEmptyc___rarg(obj*, obj*, obj*);
+obj* l_Lean_SMap_stageSizes___rarg(obj*);
 obj* l_Lean_SMap_foldStage2(obj*, obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_switch___rarg___boxed(obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_contains___main___rarg___boxed(obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_foldStage2___rarg(obj*, obj*, obj*);
 obj* l_Lean_SMap_switch___boxed(obj*, obj*);
 obj* l_Lean_SMap_empty___rarg___boxed(obj*, obj*, obj*);
+obj* l_Lean_SMap_size___boxed(obj*, obj*, obj*, obj*, obj*);
 obj* l_Lean_SMap_empty___rarg(obj* x_0, obj* x_1, obj* x_2) {
 _start:
 {
@@ -536,6 +547,92 @@ lean::dec(x_5);
 return x_6;
 }
 }
+obj* l_Lean_SMap_size___rarg(obj* x_0) {
+_start:
+{
+obj* x_1; obj* x_2; obj* x_3; obj* x_4; obj* x_5; obj* x_6; 
+x_1 = lean::cnstr_get(x_0, 0);
+x_2 = lean::cnstr_get(x_0, 1);
+x_3 = lean::mk_nat_obj(0ul);
+x_4 = l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(x_3, x_2);
+x_5 = lean::cnstr_get(x_1, 0);
+x_6 = lean::nat_add(x_5, x_4);
+lean::dec(x_4);
+return x_6;
+}
+}
+obj* l_Lean_SMap_size(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4) {
+_start:
+{
+obj* x_5; 
+x_5 = lean::alloc_closure(reinterpret_cast<void*>(l_Lean_SMap_size___rarg___boxed), 1, 0);
+return x_5;
+}
+}
+obj* l_Lean_SMap_size___rarg___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_Lean_SMap_size___rarg(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
+obj* l_Lean_SMap_size___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4) {
+_start:
+{
+obj* x_5; 
+x_5 = l_Lean_SMap_size(x_0, x_1, x_2, x_3, x_4);
+lean::dec(x_0);
+lean::dec(x_1);
+lean::dec(x_2);
+lean::dec(x_3);
+lean::dec(x_4);
+return x_5;
+}
+}
+obj* l_Lean_SMap_stageSizes___rarg(obj* x_0) {
+_start:
+{
+obj* x_1; obj* x_3; obj* x_6; obj* x_7; obj* x_9; obj* x_12; 
+x_1 = lean::cnstr_get(x_0, 0);
+lean::inc(x_1);
+x_3 = lean::cnstr_get(x_0, 1);
+lean::inc(x_3);
+lean::dec(x_0);
+x_6 = lean::mk_nat_obj(0ul);
+x_7 = l_RBNode_fold___main___at_RBMap_size___spec__1___rarg(x_6, x_3);
+lean::dec(x_3);
+x_9 = lean::cnstr_get(x_1, 0);
+lean::inc(x_9);
+lean::dec(x_1);
+x_12 = lean::alloc_cnstr(0, 2, 0);
+lean::cnstr_set(x_12, 0, x_9);
+lean::cnstr_set(x_12, 1, x_7);
+return x_12;
+}
+}
+obj* l_Lean_SMap_stageSizes(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4) {
+_start:
+{
+obj* x_5; 
+x_5 = lean::alloc_closure(reinterpret_cast<void*>(l_Lean_SMap_stageSizes___rarg), 1, 0);
+return x_5;
+}
+}
+obj* l_Lean_SMap_stageSizes___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4) {
+_start:
+{
+obj* x_5; 
+x_5 = l_Lean_SMap_stageSizes(x_0, x_1, x_2, x_3, x_4);
+lean::dec(x_0);
+lean::dec(x_1);
+lean::dec(x_2);
+lean::dec(x_3);
+lean::dec(x_4);
+return x_5;
+}
+}
 obj* initialize_init_data_hashmap_default(obj*);
 obj* initialize_init_data_rbmap_default(obj*);
 static bool _G_initialized = false;