7 files changed, 101 insertions, 13 deletions

diff --git a/crates/ra_cli/Cargo.toml b/crates/ra_cli/Cargo.toml
index b10873cd8..639b7c949 100644
--- a/crates/ra_cli/Cargo.toml
+++ b/crates/ra_cli/Cargo.toml
@@ -20,3 +20,4 @@ ra_db = { path = "../ra_db" }
 [dependencies.ra_prof]
 path = "../ra_prof"
 # features = [ "cpuprofiler" ]
+# features = [ "jemalloc" ]
diff --git a/crates/ra_cli/src/analysis_stats.rs b/crates/ra_cli/src/analysis_stats.rs
index 2eaf784a6..a01c07c3f 100644
--- a/crates/ra_cli/src/analysis_stats.rs
+++ b/crates/ra_cli/src/analysis_stats.rs
@@ -6,9 +6,9 @@ use ra_syntax::AstNode;
 
 use crate::Result;
 
-pub fn run(verbose: bool, path: &Path, only: Option<&str>) -> Result<()> {
+pub fn run(verbose: bool, memory_usage: bool, path: &Path, only: Option<&str>) -> Result<()> {
     let db_load_time = Instant::now();
-    let (host, roots) = ra_batch::load_cargo(path)?;
+    let (mut host, roots) = ra_batch::load_cargo(path)?;
     let db = host.raw_database();
     println!("Database loaded, {} roots, {:?}", roots.len(), db_load_time.elapsed());
     let analysis_time = Instant::now();
@@ -113,5 +113,12 @@ pub fn run(verbose: bool, path: &Path, only: Option<&str>) -> Result<()> {
         (num_exprs_partially_unknown * 100 / num_exprs)
     );
     println!("Analysis: {:?}, {}", analysis_time.elapsed(), ra_prof::memory_usage());
+
+    if memory_usage {
+        for (name, bytes) in host.per_query_memory_usage() {
+            println!("{:>8} {}", bytes, name)
+        }
+    }
+
     Ok(())
 }
diff --git a/crates/ra_cli/src/main.rs b/crates/ra_cli/src/main.rs
index 5adf8b096..d2f6544f8 100644
--- a/crates/ra_cli/src/main.rs
+++ b/crates/ra_cli/src/main.rs
@@ -24,6 +24,7 @@ fn main() -> Result<()> {
         .subcommand(
             SubCommand::with_name("analysis-stats")
                 .arg(Arg::with_name("verbose").short("v").long("verbose"))
+                .arg(Arg::with_name("memory-usage").long("memory-usage"))
                 .arg(Arg::with_name("only").short("o").takes_value(true))
                 .arg(Arg::with_name("path")),
         )
@@ -71,9 +72,10 @@ fn main() -> Result<()> {
         }
         ("analysis-stats", Some(matches)) => {
             let verbose = matches.is_present("verbose");
+            let memory_usage = matches.is_present("memory-usage");
             let path = matches.value_of("path").unwrap_or("");
             let only = matches.value_of("only");
-            analysis_stats::run(verbose, path.as_ref(), only)?;
+            analysis_stats::run(verbose, memory_usage, path.as_ref(), only)?;
         }
         ("analysis-bench", Some(matches)) => {
             let verbose = matches.is_present("verbose");
diff --git a/crates/ra_hir/src/db.rs b/crates/ra_hir/src/db.rs
index 8f4de1c85..b7d5ab5c6 100644
--- a/crates/ra_hir/src/db.rs
+++ b/crates/ra_hir/src/db.rs
@@ -57,17 +57,21 @@ pub trait InternDatabase: SourceDatabase {
 pub trait AstDatabase: InternDatabase {
     #[salsa::invoke(crate::source_id::AstIdMap::ast_id_map_query)]
     fn ast_id_map(&self, file_id: HirFileId) -> Arc<AstIdMap>;
+
     #[salsa::transparent]
     #[salsa::invoke(crate::source_id::AstIdMap::file_item_query)]
     fn ast_id_to_node(&self, file_id: HirFileId, ast_id: ErasedFileAstId) -> TreeArc<SyntaxNode>;
+
     #[salsa::transparent]
     #[salsa::invoke(crate::ids::HirFileId::parse_or_expand_query)]
     fn parse_or_expand(&self, file_id: HirFileId) -> Option<TreeArc<SyntaxNode>>;
+
     #[salsa::invoke(crate::ids::HirFileId::parse_macro_query)]
     fn parse_macro(&self, macro_file: ids::MacroFile) -> Option<TreeArc<SyntaxNode>>;
 
     #[salsa::invoke(crate::ids::macro_def_query)]
     fn macro_def(&self, macro_id: MacroDefId) -> Option<Arc<mbe::MacroRules>>;
+
     #[salsa::invoke(crate::ids::macro_arg_query)]
     fn macro_arg(&self, macro_call: ids::MacroCallId) -> Option<Arc<tt::Subtree>>;
 
diff --git a/crates/ra_ide_api/src/change.rs b/crates/ra_ide_api/src/change.rs
index 8d9918d16..2dfedad01 100644
--- a/crates/ra_ide_api/src/change.rs
+++ b/crates/ra_ide_api/src/change.rs
@@ -9,7 +9,7 @@ use ra_db::{
     salsa::{Database, SweepStrategy},
 };
 use ra_syntax::SourceFile;
-use ra_prof::profile;
+use ra_prof::{profile, Bytes, memory_usage};
 use relative_path::RelativePathBuf;
 use rayon::prelude::*;
 
@@ -243,4 +243,65 @@ impl RootDatabase {
         self.query(hir::db::InferQuery).sweep(sweep);
         self.query(hir::db::BodyHirQuery).sweep(sweep);
     }
+
+    pub(crate) fn per_query_memory_usage(&mut self) -> Vec<(String, Bytes)> {
+        let mut acc: Vec<(String, Bytes)> = vec![];
+        let sweep = SweepStrategy::default().discard_values().sweep_all_revisions();
+        macro_rules! sweep_each_query {
+            ($($q:path)*) => {$(
+                let before = memory_usage().allocated;
+                self.query($q).sweep(sweep);
+                let after = memory_usage().allocated;
+                let q: $q = Default::default();
+                let name = format!("{:?}", q);
+                acc.push((name, before - after));
+            )*}
+        }
+        sweep_each_query![
+            ra_db::ParseQuery
+            ra_db::SourceRootCratesQuery
+            hir::db::AstIdMapQuery
+            hir::db::ParseMacroQuery
+            hir::db::MacroDefQuery
+            hir::db::MacroArgQuery
+            hir::db::MacroExpandQuery
+            hir::db::StructDataQuery
+            hir::db::EnumDataQuery
+            hir::db::TraitDataQuery
+            hir::db::TraitItemsIndexQuery
+            hir::db::RawItemsQuery
+            hir::db::RawItemsWithSourceMapQuery
+            hir::db::CrateDefMapQuery
+            hir::db::ImplsInModuleQuery
+            hir::db::ImplsInModuleWithSourceMapQuery
+            hir::db::GenericParamsQuery
+            hir::db::FnDataQuery
+            hir::db::TypeAliasDataQuery
+            hir::db::ConstDataQuery
+            hir::db::StaticDataQuery
+            hir::db::ModuleLangItemsQuery
+            hir::db::LangItemsQuery
+            hir::db::LangItemQuery
+            hir::db::DocumentationQuery
+            hir::db::ExprScopesQuery
+            hir::db::InferQuery
+            hir::db::TypeForDefQuery
+            hir::db::TypeForFieldQuery
+            hir::db::CallableItemSignatureQuery
+            hir::db::GenericPredicatesQuery
+            hir::db::GenericDefaultsQuery
+            hir::db::BodyWithSourceMapQuery
+            hir::db::BodyHirQuery
+            hir::db::ImplsInCrateQuery
+            hir::db::ImplsForTraitQuery
+            hir::db::AssociatedTyDataQuery
+            hir::db::TraitDatumQuery
+            hir::db::StructDatumQuery
+            hir::db::ImplDatumQuery
+            hir::db::ImplementsQuery
+            hir::db::NormalizeQuery
+        ];
+        acc.sort_by_key(|it| std::cmp::Reverse(it.1));
+        acc
+    }
 }
diff --git a/crates/ra_ide_api/src/lib.rs b/crates/ra_ide_api/src/lib.rs
index 28a74c003..817e65df0 100644
--- a/crates/ra_ide_api/src/lib.rs
+++ b/crates/ra_ide_api/src/lib.rs
@@ -270,6 +270,10 @@ impl AnalysisHost {
     pub fn collect_garbage(&mut self) {
         self.db.collect_garbage();
     }
+    /// NB: this clears the database
+    pub fn per_query_memory_usage(&mut self) -> Vec<(String, ra_prof::Bytes)> {
+        self.db.per_query_memory_usage()
+    }
     pub fn raw_database(&self) -> &(impl hir::db::HirDatabase + salsa::Database) {
         &self.db
     }
diff --git a/crates/ra_prof/src/memory_usage.rs b/crates/ra_prof/src/memory_usage.rs
index 2bde8fb5f..7f857942d 100644
--- a/crates/ra_prof/src/memory_usage.rs
+++ b/crates/ra_prof/src/memory_usage.rs
@@ -27,21 +27,23 @@ impl fmt::Display for MemoryUsage {
     }
 }
 
-#[derive(Default)]
+#[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
 pub struct Bytes(usize);
 
 impl fmt::Display for Bytes {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let bytes = self.0;
-        if bytes < 4096 {
-            return write!(f, "{} bytes", bytes);
+        let mut value = bytes;
+        let mut suffix = "b";
+        if value > 4096 {
+            value /= 1024;
+            suffix = "kb";
+            if value > 4096 {
+                value /= 1024;
+                suffix = "mb";
+            }
         }
-        let kb = bytes / 1024;
-        if kb < 4096 {
-            return write!(f, "{}kb", kb);
-        }
-        let mb = kb / 1024;
-        write!(f, "{}mb", mb)
+        f.pad(&format!("{}{}", value, suffix))
     }
 }
 
@@ -50,3 +52,10 @@ impl std::ops::AddAssign<usize> for Bytes {
         self.0 += x;
     }
 }
+
+impl std::ops::Sub for Bytes {
+    type Output = Bytes;
+    fn sub(self, rhs: Bytes) -> Bytes {
+        Bytes(self.0 - rhs.0)
+    }
+}