Disassemble to YARV

Author: kddnewton

lib/syntax_tree.rb

@@ -72,6 +72,10 @@ def initialize(start_line:, start_char:, end_line:, end_char:)
       @end_char = end_char
     end
 
+    def lines
+      start_line..end_line
+    end
+
     def ==(other)
       other.is_a?(Location) && start_line == other.start_line &&
         start_char == other.start_char && end_line == other.end_line &&
@@ -3098,12 +3102,14 @@ def to_json(*opts)
   #     (nil | Ensure) ensure_clause
   #   ) -> BodyStmt
   def on_bodystmt(statements, rescue_clause, else_clause, ensure_clause)
+    ending = [ensure_clause, else_clause, rescue_clause, statements].detect(&:itself)
+
     BodyStmt.new(
       statements: statements,
       rescue_clause: rescue_clause,
       else_clause: else_clause,
       ensure_clause: ensure_clause,
-      location: Location.fixed(line: lineno, char: char_pos)
+      location: statements.location.to(ending.location)
     )
   end
 

lib/syntax_tree/code_actions.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+class SyntaxTree
+  class CodeActions
+    class DisasmAction
+      attr_reader :line, :node
+
+      def initialize(line, node)
+        @line = line
+        @node = node
+      end
+
+      def as_json
+        {
+          title: "Disasm #{node.name.value}",
+          command: "syntaxTree.disasm",
+          arguments: [line, node.name.value]
+        }
+      end
+    end
+
+    attr_reader :line, :actions
+
+    def initialize(line)
+      @line = line
+      @actions = []
+    end
+
+    def disasm(node)
+      @actions << DisasmAction.new(line, node)
+    end
+
+    def self.find(program, line)
+      code_actions = new(line)
+      queue = [program]
+
+      until queue.empty?
+        node = queue.shift
+
+        # If we're found a method definition that starts on the given line, then
+        # we're going to add the disasm code action for that node.
+        if [Def, DefEndless, Defs].include?(node.class) && node.location.start_line == line
+          code_actions.disasm(node)
+        end
+
+        # Check if the node covers the given line. If it doesn't, then we just
+        # bail out and go to the next node in the queue.
+        next unless node.location.lines.cover?(line)
+
+        # Get a list of child nodes and binary search over them to find the
+        # first child node that covers the given line. It's possible that there
+        # are multiple, but at least we'll have a handle on one of them.
+        child_nodes = node.child_nodes.compact
+        index =
+          child_nodes.bsearch_index do |child|
+            if child.location.lines.cover?(line)
+              0
+            else
+              line - child.location.start_line
+            end
+          end
+
+        # If no valid child was found, then just continue on to the next node in
+        # the queue.
+        next unless index
+
+        # First, we're going to go backward in our list until we find the start
+        # of the subset of child nodes that cover the given line. We're going to
+        # add each one to the queue.
+        current = index - 1
+        while current > 0 && child_nodes[current].location.lines.cover?(line)
+          queue << child_nodes[current]
+          current -= 1
+        end
+
+        # Next, we're going to go forward in our list until we find the end of
+        # the subset of child nodes that cover the given line. We'll add each of
+        # those to the queue as well.
+        current = index
+        while current < child_nodes.length && child_nodes[current].location.lines.cover?(line)
+          queue << child_nodes[current]
+          current += 1
+        end
+      end
+
+      code_actions
+    end
+  end
+end

lib/syntax_tree/language_server.rb

@@ -4,6 +4,7 @@
 require "json"
 require "uri"
 
+require_relative "code_actions"
 require_relative "implicits"
 
 class SyntaxTree
@@ -22,8 +23,8 @@ def run
         end
 
       while headers = input.gets("\r\n\r\n")
-        length = input.read(headers[/Content-Length: (\d+)/i, 1].to_i)
-        request = JSON.parse(length, symbolize_names: true)
+        source = input.read(headers[/Content-Length: (\d+)/i, 1].to_i)
+        request = JSON.parse(source, symbolize_names: true)
 
         case request
         in { method: "initialize", id: }
@@ -34,6 +35,8 @@ def run
         in { method: "shutdown" }
           store.clear
           return
+        in { method: "textDocument/codeAction", id:, params: { textDocument: { uri: }, range: { start: { line: } } } }
+          write(id: id, result: code_actions(store[uri], line + 1))
         in { method: "textDocument/didChange", params: { textDocument: { uri: }, contentChanges: [{ text: }, *] } }
           store[uri] = text
         in { method: "textDocument/didOpen", params: { textDocument: { uri:, text: } } }
@@ -42,6 +45,8 @@ def run
           store.delete(uri)
         in { method: "textDocument/formatting", id:, params: { textDocument: { uri: } } }
           write(id: id, result: [format(store[uri])])
+        in { method: "syntaxTree/disasm", id:, params: { textDocument: { uri:, query: { line:, name: } } } }
+          write(id: id, result: disasm(store[uri], line.to_i, name))
         in { method: "syntaxTree/implicits", id:, params: { textDocument: { uri: } } }
           write(id: id, result: implicits(store[uri]))
         in { method: "syntaxTree/visualizing", id:, params: { textDocument: { uri: } } }
@@ -60,11 +65,43 @@ def run
 
     def capabilities
       {
+        codeActionProvider: { codeActionsKinds: ["disasm"] },
         documentFormattingProvider: true,
-        textDocumentSync: { change: 1, openClose: true },
+        textDocumentSync: { change: 1, openClose: true }
       }
     end
 
+    def code_actions(source, line)
+      actions = CodeActions.find(SyntaxTree.parse(source), line).actions
+      log("Found #{actions.length} actions on line #{line}")
+
+      actions.map(&:as_json)
+    end
+
+    def disasm(source, line, name)
+      actions = CodeActions.find(SyntaxTree.parse(source), line).actions
+      log("Disassembling #{name.inspect} on line #{line.inspect}")
+
+      matched = actions.detect { |action| action.is_a?(CodeActions::DisasmAction) && action.node.name.value == name }
+      return "Unable to find method: #{name}" unless matched
+
+      # First, get an instruction sequence that encompasses the method that
+      # we're going to disassemble. It will include the method declaration,
+      # which will be the top instruction sequence.
+      location = matched.node.location
+      iseq = RubyVM::InstructionSequence.new(source[location.start_char...location.end_char])
+
+      # Next, get the first child. We do this because the parent instruction
+      # sequence is the method declaration, whereas the first child is the body
+      # of the method, which is what we're interested in.
+      method = nil
+      iseq.each_child { |child| method = child }
+
+      # Finally, return the disassembly as a string to the server, which will
+      # serialize it to JSON and return it back to the client.
+      method.disasm
+    end
+
     def format(source)
       {
         range: {
@@ -75,6 +112,10 @@ def format(source)
       }
     end
 
+    def log(message)
+      write(method: "window/logMessage", params: { type: 4, message: message })
+    end
+
     def implicits(source)
       implicits = Implicits.find(SyntaxTree.parse(source))
       serialize = ->(position, text) { { position: position, text: text } }

vscode/package.json

@@ -22,7 +22,8 @@
     "onCommand:syntaxTree.stop",
     "onCommand:syntaxTree.restart",
     "onCommand:syntaxTree.showOutputChannel",
-    "onCommand:syntaxTree.visualize"
+    "onCommand:syntaxTree.visualize",
+    "onCommand:syntaxTree.disasm"
   ],
   "main": "./out/extension",
   "contributes": {
@@ -46,6 +47,10 @@
       {
         "command": "syntaxTree.visualize",
         "title": "Syntax Tree: Visualize"
+      },
+      {
+        "command": "syntaxTree.disasm",
+        "title": "Syntax Tree: Disassemble"
       }
     ],
     "colors": [

vscode/src/Disasm.ts

@@ -0,0 +1,55 @@
+import { commands, Disposable, languages, OutputChannel, ProviderResult, TextDocumentContentProvider, Uri, ViewColumn, window, workspace } from "vscode";
+import { LanguageClient } from "vscode-languageclient/node";
+
+class Disasm implements Disposable, TextDocumentContentProvider {
+  // The client used to communicate with the language server.
+  private readonly languageClient: LanguageClient;
+
+  // The output channel used for logging for this class. It's given from the
+  // main file so that it uses the same as the rest of the extension.
+  private readonly outputChannel: OutputChannel;
+
+  // The list of callbacks and objects that should be disposed when an instance
+  // of Visualize is being disposed.
+  private readonly disposables: Disposable[];
+
+  constructor(languageClient: LanguageClient, outputChannel: OutputChannel) {
+    this.languageClient = languageClient;
+    this.outputChannel = outputChannel;
+    this.disposables = [
+      commands.registerCommand("syntaxTree.disasm", this.disasm),
+      workspace.registerTextDocumentContentProvider("syntaxTree.disasm", this)
+    ];
+  }
+
+  dispose() {
+    this.disposables.forEach((disposable) => disposable.dispose());
+  }
+
+  provideTextDocumentContent(uri: Uri): ProviderResult<string> {
+    this.outputChannel.appendLine("Requesting disassembly");
+
+    const query: Record<string, string> = {};
+    uri.query.split("&").forEach((pair) => {
+      const [key, value] = pair.split("=");
+      query[key] = value;
+    });
+
+    return this.languageClient.sendRequest("syntaxTree/disasm", { textDocument: { uri: uri.path, query } });
+  }
+
+  async disasm(line: number, name: string) {
+    const document = window.activeTextEditor?.document;
+
+    if (document && document.languageId === "ruby" && document.uri.scheme === "file") {
+      const uri = Uri.parse(`syntaxTree.disasm:${document.uri.toString()}?line=${line}&name=${name}`);
+
+      const doc = await workspace.openTextDocument(uri);
+      languages.setTextDocumentLanguage(doc, "plaintext");
+
+			await window.showTextDocument(doc, ViewColumn.Beside, true);
+    }
+  }
+}
+
+export default Disasm;

vscode/src/Visualize.ts

@@ -18,7 +18,7 @@ class Visualize implements Disposable, TextDocumentContentProvider {
     this.outputChannel = outputChannel;
     this.disposables = [
       commands.registerCommand("syntaxTree.visualize", this.visualize),
-      workspace.registerTextDocumentContentProvider("syntaxTree", this)
+      workspace.registerTextDocumentContentProvider("syntaxTree.visualize", this)
     ];
   }
 
@@ -35,7 +35,7 @@ class Visualize implements Disposable, TextDocumentContentProvider {
     const document = window.activeTextEditor?.document;
 
     if (document && document.languageId === "ruby" && document.uri.scheme === "file") {
-      const uri = Uri.parse(`syntaxTree:${document.uri.toString()}`);
+      const uri = Uri.parse(`syntaxTree.visualize:${document.uri.toString()}`);
 
       const doc = await workspace.openTextDocument(uri);
       languages.setTextDocumentLanguage(doc, "plaintext");

vscode/src/extension.ts

@@ -3,6 +3,7 @@
 import { ExtensionContext, commands, window } from "vscode";
 import { LanguageClient } from "vscode-languageclient/node";
 
+import Disasm from "./Disasm";
 import Implicits from "./Implicits";
 import startLanguageClient from "./startLanguageClient";
 import Visualize from "./Visualize";
@@ -30,6 +31,7 @@ export function activate(context: ExtensionContext) {
       await languageClient.onReady();
 
       context.subscriptions.push(
+        new Disasm(languageClient, outputChannel),
         new Implicits(languageClient, outputChannel),
         new Visualize(languageClient, outputChannel)
       );