Add a command to compute memory usage statistics

8 files changed, 74 insertions, 1 deletions

diff --git a/crates/ra_ide_db/src/change.rs b/crates/ra_ide_db/src/change.rs
index d8da3f949..84c6f40ff 100644
--- a/crates/ra_ide_db/src/change.rs
+++ b/crates/ra_ide_db/src/change.rs
@@ -164,6 +164,15 @@ impl RootDatabase {
         hir::db::BodyQuery.in_db(self).sweep(sweep);
     }
 
+    // Feature: Memory Usage
+    //
+    // Clears rust-analyzer's internal database and prints memory usage statistics.
+    //
+    // |===
+    // | Editor  | Action Name
+    //
+    // | VS Code | **Rust Analyzer: Memory Usage (Clears Database)**
+    // |===
     pub 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();
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index cc83a1406..f374334fe 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -32,7 +32,7 @@ use crate::{
     cargo_target_spec::CargoTargetSpec,
     config::RustfmtConfig,
     from_json, from_proto,
-    global_state::GlobalStateSnapshot,
+    global_state::{GlobalState, GlobalStateSnapshot},
     lsp_ext::{self, InlayHint, InlayHintsParams},
     to_proto, LspError, Result,
 };
@@ -62,6 +62,17 @@ pub(crate) fn handle_analyzer_status(snap: GlobalStateSnapshot, _: ()) -> Result
     Ok(buf)
 }
 
+pub(crate) fn handle_memory_usage(state: &mut GlobalState, _: ()) -> Result<String> {
+    let _p = profile("handle_memory_usage");
+    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);
+    }
+    Ok(out)
+}
+
 pub(crate) fn handle_syntax_tree(
     snap: GlobalStateSnapshot,
     params: lsp_ext::SyntaxTreeParams,
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index d225ad5ff..ba8a0231f 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -14,6 +14,14 @@ impl Request for AnalyzerStatus {
     const METHOD: &'static str = "rust-analyzer/analyzerStatus";
 }
 
+pub enum MemoryUsage {}
+
+impl Request for MemoryUsage {
+    type Params = ();
+    type Result = String;
+    const METHOD: &'static str = "rust-analyzer/memoryUsage";
+}
+
 pub enum ReloadWorkspace {}
 
 impl Request for ReloadWorkspace {
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index cfde55431..5900886e7 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -341,6 +341,7 @@ impl GlobalState {
             .on_sync::<lsp_ext::MatchingBrace>(|s, p| {
                 handlers::handle_matching_brace(s.snapshot(), p)
             })?
+            .on_sync::<lsp_ext::MemoryUsage>(|s, p| handlers::handle_memory_usage(s, p))?
             .on::<lsp_ext::AnalyzerStatus>(handlers::handle_analyzer_status)?
             .on::<lsp_ext::SyntaxTree>(handlers::handle_syntax_tree)?
             .on::<lsp_ext::ExpandMacro>(handlers::handle_expand_macro)?
diff --git a/editors/code/package.json b/editors/code/package.json
index 4b47fc9d3..743a2290c 100644
--- a/editors/code/package.json
+++ b/editors/code/package.json
@@ -61,6 +61,7 @@
     "activationEvents": [
         "onLanguage:rust",
         "onCommand:rust-analyzer.analyzerStatus",
+        "onCommand:rust-analyzer.memoryUsage",
         "onCommand:rust-analyzer.reloadWorkspace",
         "workspaceContains:**/Cargo.toml"
     ],
@@ -143,6 +144,11 @@
                 "category": "Rust Analyzer"
             },
             {
+                "command": "rust-analyzer.memoryUsage",
+                "title": "Memory Usage (Clears Database)",
+                "category": "Rust Analyzer"
+            },
+            {
                 "command": "rust-analyzer.reloadWorkspace",
                 "title": "Reload workspace",
                 "category": "Rust Analyzer"
@@ -844,6 +850,10 @@
                     "when": "inRustProject"
                 },
                 {
+                    "command": "rust-analyzer.memoryUsage",
+                    "when": "inRustProject"
+                },
+                {
                     "command": "rust-analyzer.reloadWorkspace",
                     "when": "inRustProject"
                 },
diff --git a/editors/code/src/commands.ts b/editors/code/src/commands.ts
index 19a9c2a0d..1f3a7cf7e 100644
--- a/editors/code/src/commands.ts
+++ b/editors/code/src/commands.ts
@@ -55,6 +55,38 @@ export function analyzerStatus(ctx: Ctx): Cmd {
     };
 }
 
+export function memoryUsage(ctx: Ctx): Cmd {
+    const tdcp = new class implements vscode.TextDocumentContentProvider {
+        readonly uri = vscode.Uri.parse('rust-analyzer-memory://memory');
+        readonly eventEmitter = new vscode.EventEmitter<vscode.Uri>();
+
+        provideTextDocumentContent(_uri: vscode.Uri): vscode.ProviderResult<string> {
+            if (!vscode.window.activeTextEditor) return '';
+
+            return ctx.client.sendRequest(ra.memoryUsage, null).then((mem) => {
+                return 'Per-query memory usage:\n' + mem + '\n(note: database has been cleared)';
+            });
+        }
+
+        get onDidChange(): vscode.Event<vscode.Uri> {
+            return this.eventEmitter.event;
+        }
+    }();
+
+    ctx.pushCleanup(
+        vscode.workspace.registerTextDocumentContentProvider(
+            'rust-analyzer-memory',
+            tdcp,
+        ),
+    );
+
+    return async () => {
+        tdcp.eventEmitter.fire(tdcp.uri);
+        const document = await vscode.workspace.openTextDocument(tdcp.uri);
+        return vscode.window.showTextDocument(document, vscode.ViewColumn.Two, true);
+    };
+}
+
 export function matchingBrace(ctx: Ctx): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
diff --git a/editors/code/src/lsp_ext.ts b/editors/code/src/lsp_ext.ts
index bf4703239..5f32cb40e 100644
--- a/editors/code/src/lsp_ext.ts
+++ b/editors/code/src/lsp_ext.ts
@@ -5,6 +5,7 @@
 import * as lc from "vscode-languageclient";
 
 export const analyzerStatus = new lc.RequestType<null, string, void>("rust-analyzer/analyzerStatus");
+export const memoryUsage = new lc.RequestType<null, string, void>("rust-analyzer/memoryUsage");
 
 export type Status = "loading" | "ready" | "invalid" | "needsReload";
 export const status = new lc.NotificationType<Status>("rust-analyzer/status");
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index 4b990afa1..eda95ae5c 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -96,6 +96,7 @@ async function tryActivate(context: vscode.ExtensionContext) {
     });
 
     ctx.registerCommand('analyzerStatus', commands.analyzerStatus);
+    ctx.registerCommand('memoryUsage', commands.memoryUsage);
     ctx.registerCommand('reloadWorkspace', commands.reloadWorkspace);
     ctx.registerCommand('matchingBrace', commands.matchingBrace);
     ctx.registerCommand('joinLines', commands.joinLines);