Count query entries in memory usage command

Author: Veykril

crates/hir-ty/src/display.rs

@@ -553,7 +553,17 @@ fn render_const_scalar(
                 render_const_scalar(f, bytes, memory_map, t)
             }
             _ => {
-                let addr = usize::from_le_bytes(b.try_into().unwrap());
+                let addr = usize::from_le_bytes(match b.try_into() {
+                    Ok(b) => b,
+                    Err(_) => {
+                        never!(
+                            "tried rendering ty {:?} in const ref with incorrect byte count {}",
+                            t,
+                            b.len()
+                        );
+                        return f.write_str("<layout-error>");
+                    }
+                });
                 let Ok(layout) = f.db.layout_of_ty(t.clone(), krate) else {
                     return f.write_str("<layout-error>");
                 };

crates/ide-db/src/apply_change.rs

@@ -1,7 +1,10 @@
 //! Applies changes to the IDE state transactionally.
 
 use base_db::{
-    salsa::{Database, Durability},
+    salsa::{
+        debug::{DebugQueryTable, TableEntry},
+        Database, Durability, Query, QueryTable,
+    },
     Change, SourceRootId,
 };
 use profile::{memory_usage, Bytes};
@@ -47,16 +50,37 @@ impl RootDatabase {
     // | VS Code | **rust-analyzer: Memory Usage (Clears Database)**
     // |===
     // image::https://user-images.githubusercontent.com/48062697/113065592-08559f00-91b1-11eb-8c96-64b88068ec02.gif[]
-    pub fn per_query_memory_usage(&mut self) -> Vec<(String, Bytes)> {
-        let mut acc: Vec<(String, Bytes)> = vec![];
+    pub fn per_query_memory_usage(&mut self) -> Vec<(String, Bytes, usize)> {
+        let mut acc: Vec<(String, Bytes, usize)> = vec![];
+
+        fn collect_query_count<'q, Q>(table: &QueryTable<'q, Q>) -> usize
+        where
+            QueryTable<'q, Q>: DebugQueryTable,
+            Q: Query,
+            <Q as Query>::Storage: 'q,
+        {
+            struct EntryCounter(usize);
+            impl<K, V> FromIterator<TableEntry<K, V>> for EntryCounter {
+                fn from_iter<T>(iter: T) -> EntryCounter
+                where
+                    T: IntoIterator<Item = TableEntry<K, V>>,
+                {
+                    EntryCounter(iter.into_iter().count())
+                }
+            }
+            table.entries::<EntryCounter>().0
+        }
+
         macro_rules! purge_each_query {
             ($($q:path)*) => {$(
                 let before = memory_usage().allocated;
-                $q.in_db(self).purge();
+                let table = $q.in_db(self);
+                let count = collect_query_count(&table);
+                table.purge();
                 let after = memory_usage().allocated;
                 let q: $q = Default::default();
                 let name = format!("{:?}", q);
-                acc.push((name, before - after));
+                acc.push((name, before - after, count));
             )*}
         }
         purge_each_query![

crates/ide/src/lib.rs

@@ -181,7 +181,7 @@ impl AnalysisHost {
     }
 
     /// NB: this clears the database
-    pub fn per_query_memory_usage(&mut self) -> Vec<(String, profile::Bytes)> {
+    pub fn per_query_memory_usage(&mut self) -> Vec<(String, profile::Bytes, usize)> {
         self.db.per_query_memory_usage()
     }
     pub fn request_cancellation(&mut self) {

crates/rust-analyzer/src/cli.rs

@@ -50,21 +50,24 @@ fn report_metric(metric: &str, value: u64, unit: &str) {
 }
 
 fn print_memory_usage(mut host: AnalysisHost, vfs: Vfs) {
-    let mut mem = host.per_query_memory_usage();
+    let mem = host.per_query_memory_usage();
 
     let before = profile::memory_usage();
     drop(vfs);
     let vfs = before.allocated - profile::memory_usage().allocated;
-    mem.push(("VFS".into(), vfs));
 
     let before = profile::memory_usage();
     drop(host);
-    mem.push(("Unaccounted".into(), before.allocated - profile::memory_usage().allocated));
+    let unaccounted = before.allocated - profile::memory_usage().allocated;
+    let remaining = profile::memory_usage().allocated;
 
-    mem.push(("Remaining".into(), profile::memory_usage().allocated));
-
-    for (name, bytes) in mem {
+    for (name, bytes, entries) in mem {
         // NOTE: Not a debug print, so avoid going through the `eprintln` defined above.
-        eprintln!("{bytes:>8} {name}");
+        eprintln!("{bytes:>8} {entries:>6} {name}");
     }
+    eprintln!("{vfs:>8}        VFS");
+
+    eprintln!("{unaccounted:>8}        Unaccounted");
+
+    eprintln!("{remaining:>8}        Remaining");
 }

crates/rust-analyzer/src/handlers/request.rs

@@ -115,13 +115,14 @@ pub(crate) fn handle_analyzer_status(
 
 pub(crate) fn handle_memory_usage(state: &mut GlobalState, _: ()) -> Result<String> {
     let _p = profile::span("handle_memory_usage");
-    let mut mem = state.analysis_host.per_query_memory_usage();
-    mem.push(("Remaining".into(), profile::memory_usage().allocated));
+    let mem = state.analysis_host.per_query_memory_usage();
 
     let mut out = String::new();
-    for (name, bytes) in mem {
-        format_to!(out, "{:>8} {}\n", bytes, name);
+    for (name, bytes, entries) in mem {
+        format_to!(out, "{:>8} {:>6} {}\n", bytes, entries, name);
     }
+    format_to!(out, "{:>8}        Remaining\n", profile::memory_usage().allocated);
+
     Ok(out)
 }