Auto merge of #15020 - Veykril:perf2, r=Veykril

Count query entries in memory usage command

crates/hir-def/src/lib.rs

@@ -1028,13 +1028,13 @@ fn attr_macro_as_call_id(
     def: MacroDefId,
 ) -> MacroCallId {
     let arg = match macro_attr.input.as_deref() {
-        Some(AttrInput::TokenTree(tt, map)) => (
+        Some(AttrInput::TokenTree(tt)) => (
             {
-                let mut tt = tt.clone();
+                let mut tt = tt.0.clone();
                 tt.delimiter = tt::Delimiter::UNSPECIFIED;
                 tt
             },
-            map.clone(),
+            tt.1.clone(),
         ),
         _ => (tt::Subtree::empty(), Default::default()),
     };

crates/hir-expand/src/attrs.rs

@@ -192,14 +192,14 @@ pub enum AttrInput {
     /// `#[attr = "string"]`
     Literal(SmolStr),
     /// `#[attr(subtree)]`
-    TokenTree(tt::Subtree, mbe::TokenMap),
+    TokenTree(Box<(tt::Subtree, mbe::TokenMap)>),
 }
 
 impl fmt::Display for AttrInput {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             AttrInput::Literal(lit) => write!(f, " = \"{}\"", lit.escape_debug()),
-            AttrInput::TokenTree(subtree, _) => subtree.fmt(f),
+            AttrInput::TokenTree(tt) => tt.0.fmt(f),
         }
     }
 }
@@ -220,7 +220,7 @@ impl Attr {
             Some(Interned::new(AttrInput::Literal(value)))
         } else if let Some(tt) = ast.token_tree() {
             let (tree, map) = syntax_node_to_token_tree(tt.syntax());
-            Some(Interned::new(AttrInput::TokenTree(tree, map)))
+            Some(Interned::new(AttrInput::TokenTree(Box::new((tree, map)))))
         } else {
             None
         };
@@ -256,7 +256,7 @@ impl Attr {
     /// #[path(ident)]
     pub fn single_ident_value(&self) -> Option<&tt::Ident> {
         match self.input.as_deref()? {
-            AttrInput::TokenTree(subtree, _) => match &*subtree.token_trees {
+            AttrInput::TokenTree(tt) => match &*tt.0.token_trees {
                 [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] => Some(ident),
                 _ => None,
             },
@@ -267,7 +267,7 @@ impl Attr {
     /// #[path TokenTree]
     pub fn token_tree_value(&self) -> Option<&Subtree> {
         match self.input.as_deref()? {
-            AttrInput::TokenTree(subtree, _) => Some(subtree),
+            AttrInput::TokenTree(tt) => Some(&tt.0),
             _ => None,
         }
     }

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)
 }