implement first pass of memory layout viewer

9 files changed, 561 insertions, 6 deletions

diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index cb2a1140ba8..a9a8f6903b2 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -60,6 +60,7 @@ mod interpret_function;
 mod view_item_tree;
 mod shuffle_crate_graph;
 mod fetch_crates;
+mod view_memory_layout;
 
 use std::ffi::OsStr;
 
@@ -74,6 +75,7 @@ use ide_db::{
 };
 use syntax::SourceFile;
 use triomphe::Arc;
+use view_memory_layout::{view_memory_layout, RecursiveMemoryLayout};
 
 use crate::navigation_target::{ToNav, TryToNav};
 
@@ -724,6 +726,10 @@ impl Analysis {
         self.with_db(|db| move_item::move_item(db, range, direction))
     }
 
+    pub fn get_recursive_memory_layout(&self, position: FilePosition) -> Cancellable<Option<RecursiveMemoryLayout>> {
+        self.with_db(|db| view_memory_layout(db, position))
+    }
+
     /// Performs an operation on the database that may be canceled.
     ///
     /// rust-analyzer needs to be able to answer semantic questions about the
diff --git a/crates/ide/src/view_memory_layout.rs b/crates/ide/src/view_memory_layout.rs
new file mode 100644
index 00000000000..3b1b4968e77
--- /dev/null
+++ b/crates/ide/src/view_memory_layout.rs
@@ -0,0 +1,170 @@
+use hir::{Field, HirDisplay, Layout, Semantics, Type};
+use ide_db::{
+    defs::{Definition, IdentClass},
+    helpers::pick_best_token,
+    RootDatabase,
+};
+use syntax::{AstNode, SyntaxKind, SyntaxToken};
+
+use crate::FilePosition;
+
+pub struct MemoryLayoutNode {
+    pub item_name: String,
+    pub typename: String,
+    pub size: u64,
+    pub alignment: u64,
+    pub offset: u64,
+    pub parent_idx: i64,
+    pub children_start: i64,
+    pub children_len: u64,
+}
+
+pub struct RecursiveMemoryLayout {
+    pub nodes: Vec<MemoryLayoutNode>,
+}
+
+fn get_definition(sema: &Semantics<'_, RootDatabase>, token: SyntaxToken) -> Option<Definition> {
+    for token in sema.descend_into_macros(token) {
+        let def = IdentClass::classify_token(sema, &token).map(IdentClass::definitions_no_ops);
+        if let Some(&[x]) = def.as_deref() {
+            return Some(x);
+        }
+    }
+    None
+}
+
+pub(crate) fn view_memory_layout(
+    db: &RootDatabase,
+    position: FilePosition,
+) -> Option<RecursiveMemoryLayout> {
+    let sema = Semantics::new(db);
+    let file = sema.parse(position.file_id);
+    let token =
+        pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
+            SyntaxKind::IDENT => 3,
+            _ => 0,
+        })?;
+
+    let def = get_definition(&sema, token)?;
+
+    let ty = match def {
+        Definition::Adt(it) => it.ty(db),
+        Definition::TypeAlias(it) => it.ty(db),
+        Definition::BuiltinType(it) => it.ty(db),
+        Definition::SelfType(it) => it.self_ty(db),
+        Definition::Local(it) => it.ty(db),
+        _ => return None,
+    };
+
+    enum FieldOrTupleIdx {
+        Field(Field),
+        TupleIdx(usize),
+    }
+
+    impl FieldOrTupleIdx {
+        fn name(&self, db: &RootDatabase) -> String {
+            match *self {
+                FieldOrTupleIdx::Field(f) => f
+                    .name(db)
+                    .as_str()
+                    .map(|s| s.to_owned())
+                    .unwrap_or_else(|| format!("{:#?}", f.name(db).as_tuple_index().unwrap())),
+                FieldOrTupleIdx::TupleIdx(i) => format!(".{i}").to_owned(),
+            }
+        }
+
+        fn index(&self) -> usize {
+            match *self {
+                FieldOrTupleIdx::Field(f) => f.index(),
+                FieldOrTupleIdx::TupleIdx(i) => i,
+            }
+        }
+    }
+
+    fn read_layout(
+        nodes: &mut Vec<MemoryLayoutNode>,
+        db: &RootDatabase,
+        ty: &Type,
+        layout: &Layout,
+        parent_idx: usize,
+    ) {
+        let mut fields = ty
+            .fields(db)
+            .into_iter()
+            .map(|(f, ty)| (FieldOrTupleIdx::Field(f), ty))
+            .chain(
+                ty.tuple_fields(db)
+                    .into_iter()
+                    .enumerate()
+                    .map(|(i, ty)| (FieldOrTupleIdx::TupleIdx(i), ty)),
+            )
+            .collect::<Vec<_>>();
+
+        fields.sort_by_key(|(f, _)| layout.field_offset(f.index()).unwrap_or(u64::MAX));
+
+        if fields.len() == 0 {
+            return;
+        }
+
+        let children_start = nodes.len();
+        nodes[parent_idx].children_start = children_start as i64;
+        nodes[parent_idx].children_len = fields.len() as u64;
+
+        for (field, child_ty) in fields.iter() {
+            if let Ok(child_layout) = child_ty.layout(db) {
+                nodes.push(MemoryLayoutNode {
+                    item_name: field.name(db),
+                    typename: child_ty.display(db).to_string(),
+                    size: child_layout.size(),
+                    alignment: child_layout.align(),
+                    offset: layout.field_offset(field.index()).unwrap_or(0),
+                    parent_idx: parent_idx as i64,
+                    children_start: -1,
+                    children_len: 0,
+                });
+            } else {
+                nodes.push(MemoryLayoutNode {
+                    item_name: field.name(db)
+                        + format!("(no layout data: {:?})", child_ty.layout(db).unwrap_err())
+                            .as_ref(),
+                    typename: child_ty.display(db).to_string(),
+                    size: 0,
+                    offset: 0,
+                    alignment: 0,
+                    parent_idx: parent_idx as i64,
+                    children_start: -1,
+                    children_len: 0,
+                });
+            }
+        }
+
+        for (i, (_, child_ty)) in fields.iter().enumerate() {
+            if let Ok(child_layout) = child_ty.layout(db) {
+                read_layout(nodes, db, &child_ty, &child_layout, children_start + i);
+            }
+        }
+    }
+
+    ty.layout(db).map(|layout| {
+        let item_name = match def {
+            Definition::Local(l) => l.name(db).as_str().unwrap().to_owned(),
+            _ => "[ROOT]".to_owned(),
+        };
+
+        let typename = ty.display(db).to_string();
+
+        let mut nodes = vec![MemoryLayoutNode {
+            item_name,
+            typename: typename.clone(),
+            size: layout.size(),
+            offset: 0,
+            alignment: layout.align(),
+            parent_idx: -1,
+            children_start: -1,
+            children_len: 0,
+        }];
+        read_layout(&mut nodes, db, &ty, &layout, 0);
+
+        RecursiveMemoryLayout { nodes }
+    })
+}
diff --git a/crates/rust-analyzer/src/handlers/request.rs b/crates/rust-analyzer/src/handlers/request.rs
index bb0c7ffa728..2b901a60191 100644
--- a/crates/rust-analyzer/src/handlers/request.rs
+++ b/crates/rust-analyzer/src/handlers/request.rs
@@ -1689,6 +1689,34 @@ pub(crate) fn handle_move_item(
     }
 }
 
+pub(crate) fn handle_view_recursive_memory_layout(
+    snap: GlobalStateSnapshot,
+    params: lsp_ext::ViewRecursiveMemoryLayoutParams,
+) -> Result<Option<lsp_ext::RecursiveMemoryLayout>> {
+    let _p = profile::span("view_recursive_memory_layout");
+    let file_id = from_proto::file_id(&snap, &params.text_document.uri)?;
+    let line_index = snap.file_line_index(file_id)?;
+    let offset = from_proto::offset(&line_index, params.position)?;
+
+    let res = snap.analysis.get_recursive_memory_layout(FilePosition { file_id, offset })?;
+    Ok(res.map(|it| lsp_ext::RecursiveMemoryLayout {
+        nodes: it
+            .nodes
+            .iter()
+            .map(|n| lsp_ext::MemoryLayoutNode {
+                item_name: n.item_name.clone(),
+                typename: n.typename.clone(),
+                size: n.size,
+                offset: n.offset,
+                alignment: n.alignment,
+                parent_idx: n.parent_idx,
+                children_start: n.children_start,
+                children_len: n.children_len,
+            })
+            .collect(),
+    }))
+}
+
 fn to_command_link(command: lsp_types::Command, tooltip: String) -> lsp_ext::CommandLink {
     lsp_ext::CommandLink { tooltip: Some(tooltip), command }
 }
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index 4d67c8b305f..39b8e840288 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -182,6 +182,40 @@ pub struct ExpandedMacro {
     pub expansion: String,
 }
 
+pub enum ViewRecursiveMemoryLayout {}
+
+impl Request for ViewRecursiveMemoryLayout {
+    type Params = ViewRecursiveMemoryLayoutParams;
+    type Result = Option<RecursiveMemoryLayout>;
+    const METHOD: &'static str = "rust-analyzer/viewRecursiveMemoryLayout";
+}
+
+#[derive(Deserialize, Serialize, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct ViewRecursiveMemoryLayoutParams {
+    pub text_document: TextDocumentIdentifier,
+    pub position: Position,
+}
+
+#[derive(Deserialize, Serialize, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct RecursiveMemoryLayout {
+    pub nodes: Vec<MemoryLayoutNode>,
+}
+
+#[derive(Deserialize, Serialize, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct MemoryLayoutNode {
+    pub item_name: String,
+    pub typename: String,
+    pub size: u64,
+    pub offset: u64,
+    pub alignment: u64,
+    pub parent_idx: i64,
+    pub children_start: i64,
+    pub children_len: u64,
+}
+
 pub enum CancelFlycheck {}
 
 impl Notification for CancelFlycheck {
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 9eecf4957af..74036710fa3 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -753,6 +753,7 @@ impl GlobalState {
             )
             .on::<lsp_types::request::WillRenameFiles>(handlers::handle_will_rename_files)
             .on::<lsp_ext::Ssr>(handlers::handle_ssr)
+            .on::<lsp_ext::ViewRecursiveMemoryLayout>(handlers::handle_view_recursive_memory_layout)
             .finish();
     }
 
diff --git a/editors/code/package.json b/editors/code/package.json
index 1a1d76c6991..cd8b40e3557 100644
--- a/editors/code/package.json
+++ b/editors/code/package.json
@@ -284,6 +284,11 @@
                 "command": "rust-analyzer.revealDependency",
                 "title": "Reveal File",
                 "category": "rust-analyzer"
+            },
+            {
+                "command": "rust-analyzer.viewMemoryLayout",
+                "title": "View Memory Layout",
+                "category": "rust-analyzer"
             }
         ],
         "keybindings": [
@@ -2067,6 +2072,10 @@
                 {
                     "command": "rust-analyzer.openCargoToml",
                     "when": "inRustProject"
+                },
+                {
+                    "command": "rust-analyzer.viewMemoryLayout",
+                    "when": "inRustProject"
                 }
             ],
             "editor/context": [
diff --git a/editors/code/src/commands.ts b/editors/code/src/commands.ts
index 3c6105e89fe..c0d68881c23 100644
--- a/editors/code/src/commands.ts
+++ b/editors/code/src/commands.ts
@@ -1129,3 +1129,286 @@ export function linkToCommand(_: Ctx): Cmd {
         }
     };
 }
+
+export function viewMemoryLayout(ctx: CtxInit): Cmd {
+    return async () => {
+
+        const editor = vscode.window.activeTextEditor;
+        if (!editor) return "";
+        const client = ctx.client;
+
+        const position = editor.selection.active;
+        const expanded = await client.sendRequest(ra.viewRecursiveMemoryLayout, {
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(
+                editor.document
+            ),
+            position,
+        });
+
+        // if (expanded == null) return "Not available";
+
+
+        const document = vscode.window.createWebviewPanel(
+            "memory_layout",
+            "[Memory Layout]",
+            vscode.ViewColumn.Two,
+            { enableScripts: true, });
+
+        document.webview.html = `<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta http-equiv="X-UA-Compatible" content="IE=edge">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Document</title>
+    <style>   
+        * {
+            box-sizing: border-box;
+        }
+
+        body {
+            margin: 0;
+            overflow: hidden;
+            min-height: 100%;
+            height: 100vh;
+            padding: 32px;
+            position: relative;
+            display: block;
+
+            background-color: var(--vscode-editor-background);
+            font-family: var(--vscode-editor-font-family);
+            font-size: var(--vscode-editor-font-size);
+            color: var(--vscode-editor-foreground);
+        }
+
+        .container {
+            position: relative;
+        }
+
+        .trans {
+            transition: all 0.2s ease-in-out;
+        }
+
+        .grid {
+            height: 100%;
+            position: relative;
+            color: var(--vscode-commandCenter-activeBorder);
+            pointer-events: none;
+        }
+
+        .grid-line {
+            position: absolute;
+            width: 100%;
+            height: 1px;
+            background-color: var(--vscode-commandCenter-activeBorder);
+        }
+
+        #tooltip {
+            position: fixed;
+            display: none;
+            z-index: 1;
+            pointer-events: none;
+            padding: 4px 8px;
+            z-index: 2;
+
+            color: var(--vscode-editorHoverWidget-foreground);
+            background-color: var(--vscode-editorHoverWidget-background);
+            border: 1px solid var(--vscode-editorHoverWidget-border);
+        }
+
+        #tooltip b {
+            color: var(--vscode-editorInlayHint-typeForeground);
+        }
+
+        #tooltip ul {
+            margin-left: 0;
+            padding-left: 20px;
+        }
+
+        table {
+            position: absolute;
+            transform: rotateZ(90deg) rotateX(180deg);
+            transform-origin: top left;
+            border-collapse: collapse;
+            table-layout: fixed;
+            left: 48px;
+            top: 0;
+            max-height: calc(100vw - 64px - 48px);
+            z-index: 1;
+        }
+
+        td {
+            border: 1px solid var(--vscode-focusBorder);
+            writing-mode: vertical-rl;
+            text-orientation: sideways-right;
+
+            height: 80px;
+        }
+
+        td p {
+            height: calc(100% - 16px);
+            width: calc(100% - 8px);
+            margin: 8px 4px;
+            display: inline-block;
+            transform: rotateY(180deg);
+            pointer-events: none;
+            overflow: hidden;
+        }
+
+        td p * {
+            overflow: hidden;
+            white-space: nowrap;
+            text-overflow: ellipsis;
+            display: inline-block;
+            height: 100%;
+        }
+
+        td p b {
+            color: var(--vscode-editorInlayHint-typeForeground);
+        }
+
+        td:hover {
+            background-color: var(--vscode-editor-hoverHighlightBackground);
+        }
+
+        td:empty {
+            visibility: hidden;
+            border: 0;
+        }
+    </style>
+</head>
+<body>
+    <div id="tooltip"></div>
+</body>
+<script>(function() {
+
+const data = ${JSON.stringify(expanded)}
+
+if (!(data && data.nodes.length)) {
+    document.body.innerText = "Not Available"
+    return
+}
+
+data.nodes.map(n => {
+    n.typename = n.typename.replaceAll('&', '&amp;').replaceAll('<', '&lt;').replaceAll('>', '&gt;').replaceAll('"', ' & quot; ').replaceAll("'", '&#039;')
+    return n
+})
+
+let height = window.innerHeight - 64
+
+addEventListener("resize", e => {
+    const new_height = window.innerHeight - 64
+    height = new_height
+    container.classList.remove("trans")
+    table.classList.remove("trans")
+    locate()
+    setTimeout(() => { // give delay to redraw, annoying but needed
+        container.classList.add("trans")
+        table.classList.add("trans")
+    }, 0)
+})
+
+const container = document.createElement("div")
+container.classList.add("container")
+container.classList.add("trans")
+document.body.appendChild(container)
+
+const tooltip = document.getElementById("tooltip")
+
+let y = 0
+let zoom = 1.0
+
+const table = document.createElement("table")
+table.classList.add("trans")
+container.appendChild(table)
+const rows = []
+
+function node_t(idx, depth, offset) {
+    if (!rows[depth]) {
+        rows[depth] = { el: document.createElement("tr"), offset: 0 }
+    }
+
+    if (rows[depth].offset < offset) {
+        const pad = document.createElement("td")
+        pad.colSpan = offset - rows[depth].offset
+        rows[depth].el.appendChild(pad)
+        rows[depth].offset += offset - rows[depth].offset
+    }
+
+    const td = document.createElement("td")
+    td.innerHTML = '<p><span>' + data.nodes[idx].itemName + ':</span> <b>' + data.nodes[idx].typename + '</b></p>'
+
+    td.colSpan = data.nodes[idx].size
+
+    td.addEventListener("mouseover", e => {
+        const node = data.nodes[idx]
+        tooltip.innerHTML = node.itemName + ": <b>" + node.typename + "</b><br/>"
+            + "<ul>"
+            + "<li>size = " + node.size + "</li>"
+            + "<li>align = " + node.alignment + "</li>"
+            + "<li>field offset = " + node.offset + "</li>"
+            + "</ul>"
+            + "<i>double click to focus</i>"
+
+        tooltip.style.display = "block"
+    })
+    td.addEventListener("mouseleave", _ => tooltip.style.display = "none")
+    const total_offset = rows[depth].offset
+    td.addEventListener("dblclick", e => {
+        const node = data.nodes[idx]
+        zoom = data.nodes[0].size / node.size
+        y = -(total_offset) / data.nodes[0].size * zoom
+        x = 0
+        locate()
+    })
+
+    rows[depth].el.appendChild(td)
+    rows[depth].offset += data.nodes[idx].size
+
+
+    if (data.nodes[idx].childrenStart != -1) {
+        for (let i = 0; i < data.nodes[idx].childrenLen; i++) {
+            if (data.nodes[data.nodes[idx].childrenStart + i].size) {
+                node_t(data.nodes[idx].childrenStart + i, depth + 1, offset + data.nodes[data.nodes[idx].childrenStart + i].offset)
+            }
+        }
+    }
+}
+
+node_t(0, 0, 0)
+
+for (const row of rows) table.appendChild(row.el)
+
+const grid = document.createElement("div")
+grid.classList.add("grid")
+container.appendChild(grid)
+
+for (let i = 0; i < data.nodes[0].size / 8 + 1; i++) {
+    const el = document.createElement("div")
+    el.classList.add("grid-line")
+    el.style.top = (i / (data.nodes[0].size / 8) * 100) + "%"
+    el.innerText = i * 8
+    grid.appendChild(el)
+}
+
+addEventListener("mousemove", e => {
+    tooltip.style.top = e.clientY + 10 + "px"
+    tooltip.style.left = e.clientX + 10 + "px"
+})
+
+function locate() {
+    container.style.top = height * y + "px"
+    container.style.height = (height * zoom) + "px"
+
+    table.style.width = container.style.height
+}
+
+locate()
+
+})()
+</script>
+</html>`
+
+        ctx.pushExtCleanup(document);
+    };
+}
diff --git a/editors/code/src/lsp_ext.ts b/editors/code/src/lsp_ext.ts
index b72804e510c..4244098ae1d 100644
--- a/editors/code/src/lsp_ext.ts
+++ b/editors/code/src/lsp_ext.ts
@@ -70,7 +70,7 @@ export const viewItemTree = new lc.RequestType<ViewItemTreeParams, string, void>
 
 export type AnalyzerStatusParams = { textDocument?: lc.TextDocumentIdentifier };
 
-export interface FetchDependencyListParams {}
+export interface FetchDependencyListParams { }
 
 export interface FetchDependencyListResult {
     crates: {
@@ -86,7 +86,7 @@ export const fetchDependencyList = new lc.RequestType<
     void
 >("rust-analyzer/fetchDependencyList");
 
-export interface FetchDependencyGraphParams {}
+export interface FetchDependencyGraphParams { }
 
 export interface FetchDependencyGraphResult {
     crates: {
@@ -150,6 +150,9 @@ export const serverStatus = new lc.NotificationType<ServerStatusParams>(
     "experimental/serverStatus"
 );
 export const ssr = new lc.RequestType<SsrParams, lc.WorkspaceEdit, void>("experimental/ssr");
+export const viewRecursiveMemoryLayout = new lc.RequestType<ViewRecursiveMemoryLayoutParams, RecursiveMemoryLayout | null, void>(
+    "rust-analyzer/viewRecursiveMemoryLayout"
+);
 
 export type JoinLinesParams = {
     textDocument: lc.TextDocumentIdentifier;
@@ -197,3 +200,23 @@ export type SsrParams = {
     position: lc.Position;
     selections: readonly lc.Range[];
 };
+
+export type ViewRecursiveMemoryLayoutParams = {
+    textDocument: lc.TextDocumentIdentifier;
+    position: lc.Position;
+};
+export type RecursiveMemoryLayoutNode = {
+    item_name: string;
+    typename: string;
+    size: number;
+    alignment: number;
+    offset: number;
+    parent_idx: number;
+    children_start: number;
+    children_len: number;
+};
+export type RecursiveMemoryLayout = {
+    name: string;
+    expansion: string;
+    nodes: RecursiveMemoryLayoutNode[];
+};
\ No newline at end of file
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index be9bc9d363c..492275968f3 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -24,11 +24,11 @@ export async function activate(
         vscode.window
             .showWarningMessage(
                 `You have both the rust-analyzer (rust-lang.rust-analyzer) and Rust (rust-lang.rust) ` +
-                    "plugins enabled. These are known to conflict and cause various functions of " +
-                    "both plugins to not work correctly. You should disable one of them.",
+                "plugins enabled. These are known to conflict and cause various functions of " +
+                "both plugins to not work correctly. You should disable one of them.",
                 "Got it"
             )
-            .then(() => {}, console.error);
+            .then(() => { }, console.error);
     }
 
     const ctx = new Ctx(context, createCommands(), fetchWorkspace());
@@ -144,7 +144,7 @@ function createCommands(): Record<string, CommandFactory> {
                     health: "stopped",
                 });
             },
-            disabled: (_) => async () => {},
+            disabled: (_) => async () => { },
         },
 
         analyzerStatus: { enabled: commands.analyzerStatus },
@@ -179,6 +179,7 @@ function createCommands(): Record<string, CommandFactory> {
         runFlycheck: { enabled: commands.runFlycheck },
         ssr: { enabled: commands.ssr },
         serverVersion: { enabled: commands.serverVersion },
+        viewMemoryLayout: { enabled: commands.viewMemoryLayout },
         // Internal commands which are invoked by the server.
         applyActionGroup: { enabled: commands.applyActionGroup },
         applySnippetWorkspaceEdit: { enabled: commands.applySnippetWorkspaceEditCommand },