feat(rustc_unpretty): use rustc inspect mir (#192)

Co-authored-by: saying121 <saying121@example.com>

CHANGELOG.md

@@ -6,6 +6,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- `:RustLsp rustcUnpretty` command:
+  Use `rustc -Z unpretty=mir` to inspect mir and other things,
+  and achieve an experience similar to Rust Playground.
+- Config: `tools.rustc_unpretty` arguments for `rustc`.
+
 ## [4.0.3] - 2024-01-28
 
 ### Fixed

README.md

@@ -491,6 +491,25 @@ vim.keymap.set(
   vim.cmd.RustLsp { 'view', 'hir' }
   vim.cmd.RustLsp { 'view', 'mir' }
   ```
+<details>
+  <summary>
+	<b>Rustc Unpretty</b>
+  </summary>
+
+  Opens a buffer with a textual representation of the MIR or others things
+  of the function containing the cursor.
+  Achieve an experience similar to Rust Playground.
+
+  NOTE: the feature need [nvim-treesitter](https://github.com/nvim-treesitter/nvim-treesitter)
+  install rust parser and need nightly rustc.
+
+  ```vimscript
+  :RustLsp rustcUnpretty [hir|mir|...]
+  ```
+  ```lua
+  vim.cmd.RustLsp { 'rustcUnpretty', 'hir' }
+  vim.cmd.RustLsp { 'rustcUnpretty', 'mir' }
+  ```
 </details>
 
 <!-- markdownlint-restore -->

lua/rustaceanvim/commands/init.lua

@@ -97,6 +97,27 @@ local command_tbl = {
     local cmd = args[1] or 'run'
     require('rustaceanvim.commands.fly_check')(cmd)
   end,
+  rustcUnpretty = function(args)
+    local err_msg = table.concat(require('rustaceanvim.commands.rustc_unpretty').available_unpretty, ' | ')
+    if not args or #args == 0 then
+      vim.notify('Expected argument list: ' .. err_msg, vim.log.levels.ERROR)
+      return
+    end
+    local arg = args[1]:lower()
+    local available = false
+    for _, value in ipairs(require('rustaceanvim.commands.rustc_unpretty').available_unpretty) do
+      if value == arg then
+        available = true
+        break
+      end
+    end
+    if not available then
+      vim.notify('Expected argument list: ' .. err_msg, vim.log.levels.ERROR)
+      return
+    end
+
+    require('rustaceanvim.commands.rustc_unpretty').rustc_unpretty(arg)
+  end,
   view = function(args)
     if not args or #args == 0 then
       vim.notify("Expected argument: 'mir' or 'hir'", vim.log.levels.ERROR)
@@ -166,6 +187,9 @@ function M.create_rust_lsp_command()
       if cmdline:match(match_start .. '%sview' .. subcmd_match) then
         return { 'mir', 'hir' }
       end
+      if cmdline:match(match_start .. '%srustcUnpretty' .. subcmd_match) then
+        return require('rustaceanvim.commands.rustc_unpretty').available_unpretty
+      end
       if cmdline:match(match_start .. '%s+%w*$') then
         return vim.tbl_filter(function(command)
           return command:find(arg_lead) ~= nil

lua/rustaceanvim/commands/rustc_unpretty.lua

@@ -0,0 +1,143 @@
+local M = {}
+
+local config = require('rustaceanvim.config.internal')
+local compat = require('rustaceanvim.compat')
+local ui = require('rustaceanvim.ui')
+local api = vim.api
+local ts = vim.treesitter
+
+local rustc = 'rustc'
+
+M.available_unpretty = {
+  'normal',
+  'identified',
+  'expanded',
+  'expanded,identified',
+  'expanded,hygiene',
+  'ast-tree',
+  'ast-tree,expanded',
+  'hir',
+  'hir,identified',
+  'hir,typed',
+  'hir-tree',
+  'thir-tree',
+  'thir-flat',
+  'mir',
+  'stable-mir',
+  'mir-cfg',
+}
+---@alias rustcir_level 'normal'| 'identified'| 'expanded'| 'expanded,identified'| 'expanded,hygiene'| 'ast-tree'| 'ast-tree,expanded'| 'hir'| 'hir,identified'| 'hir,typed'| 'hir-tree'| 'thir-tree'| 'thir-flat'| 'mir'| 'stable-mir'| 'mir-cfg'
+
+---@type integer | nil
+local latest_buf_id = nil
+
+---Get a compatible vim range (1 index based) from a TS node range.
+---
+---TS nodes start with 0 and the end col is ending exclusive.
+---They also treat a EOF/EOL char as a char ending in the first
+---col of the next row.
+---comment
+---@param range integer[]
+---@param buf integer|nil
+---@return integer, integer, integer, integer
+local function get_vim_range(range, buf)
+  ---@type integer, integer, integer, integer
+  local srow, scol, erow, ecol = unpack(range)
+  srow = srow + 1
+  scol = scol + 1
+  erow = erow + 1
+
+  if ecol == 0 then
+    -- Use the value of the last col of the previous row instead.
+    erow = erow - 1
+    if not buf or buf == 0 then
+      ecol = vim.fn.col { erow, '$' } - 1
+    else
+      ecol = #vim.api.nvim_buf_get_lines(buf, erow - 1, erow, false)[1]
+    end
+    ecol = math.max(ecol, 1)
+  end
+
+  return srow, scol, erow, ecol
+end
+
+---@param node TSNode
+local function get_rows(node)
+  local start_row, _, end_row, _ = get_vim_range({ ts.get_node_range(node) }, 0)
+  return vim.api.nvim_buf_get_lines(0, start_row - 1, end_row, true)
+end
+
+---@param sc vim.SystemCompleted
+local function handler(sc)
+  if sc.code ~= 0 then
+    vim.notify('rustc unpretty failed' .. sc.stderr, vim.log.levels.ERROR)
+    return
+  end
+
+  -- check if a buffer with the latest id is already open, if it is then
+  -- delete it and continue
+  ui.delete_buf(latest_buf_id)
+
+  -- create a new buffer
+  latest_buf_id = vim.api.nvim_create_buf(false, true) -- not listed and scratch
+
+  -- split the window to create a new buffer and set it to our window
+  ui.split(true, latest_buf_id)
+
+  local lines = vim.split(sc.stdout, '\n')
+
+  -- set filetype to rust for syntax highlighting
+  vim.bo[latest_buf_id].filetype = 'rust'
+  -- write the expansion content to the buffer
+  vim.api.nvim_buf_set_lines(latest_buf_id, 0, 0, false, lines)
+end
+
+---@param level rustcir_level
+function M.rustc_unpretty(level)
+  if #api.nvim_get_runtime_file('parser/rust.so', true) == 0 then
+    vim.notify("a treesitter parser for Rust is required for 'rustc unpretty'", vim.log.levels.ERROR)
+    return
+  end
+  if vim.fn.executable(rustc) ~= 1 then
+    vim.notify('rustc is needed to rustc unpretty.', vim.log.levels.ERROR)
+    return
+  end
+
+  local text
+
+  local cursor = api.nvim_win_get_cursor(0)
+  local pos = { math.max(cursor[1] - 1, 0), cursor[2] }
+
+  local cline = api.nvim_get_current_line()
+  if not string.find(cline, 'fn%s*') then
+    local temp = vim.fn.searchpos('fn ', 'bcn', vim.fn.line('w0'))
+    pos = { math.max(temp[1] - 1, 0), temp[2] }
+  end
+
+  local node = ts.get_node { pos = pos }
+
+  if node == nil or node:type() ~= 'function_item' then
+    vim.notify('no function found or function is incomplete', vim.log.levels.ERROR)
+    return
+  end
+
+  local b = get_rows(node)
+  if b == nil then
+    vim.notify('get code text failed', vim.log.levels.ERROR)
+    return
+  end
+  text = table.concat(b, '\n')
+
+  -- rustc need a main function for `-Z unpretty`
+  if not string.find(text, 'fn%s*main') then
+    text = text .. 'fn main() {}'
+  end
+
+  compat.system(
+    { rustc, '--edition', config.tools.rustc_unpretty.edition, '-Z', 'unpretty=' .. level, '-' },
+    { stdin = text },
+    vim.schedule_wrap(handler)
+  )
+end
+
+return M

lua/rustaceanvim/config/check.lua

@@ -65,6 +65,13 @@ function M.validate(cfg)
   if not ok then
     return false, err
   end
+  local rustc_unpretty = tools.rustc_unpretty
+  ok, err = validate('tools.rustc_ir', {
+    edition = { rustc_unpretty.edition, 'string' },
+  })
+  if not ok then
+    return false, err
+  end
   ok, err = validate('tools', {
     executor = { tools.executor, { 'table', 'string' } },
     on_initialized = { tools.on_initialized, 'function', true },

lua/rustaceanvim/config/init.lua

@@ -63,6 +63,7 @@ vim.g.rustaceanvim = vim.g.rustaceanvim
 ---@field float_win_config? table Options applied to floating windows. See |api-win_config|.
 ---@field create_graph? RustaceanCrateGraphConfig Options for showing the crate graph based on graphviz and the dot
 ---@field open_url? fun(url:string):nil If set, overrides how to open URLs
+---@field rustc_unpretty? table Options for `rustcUnpretty`
 
 ---@class RustaceanHoverActionsOpts
 ---@field replace_builtin_hover? boolean Whether to replace Neovim's built-in `vim.lsp.buf.hover` with hover actions. Default: `true`

lua/rustaceanvim/config/internal.lua

@@ -184,6 +184,12 @@ local RustaceanDefaultConfig = {
     open_url = function(url)
       require('rustaceanvim.os').open_url(url)
     end,
+    ---settings for rustc -Z unpretty
+    ---@type table
+    rustc_unpretty = {
+      ---@type string
+      edition = '2021',
+    },
   },
 
   --- all the opts to send to the LSP client