undotree.lua (11092B)
1 --- @class (private) vim.undotree.tree.entry 2 --- @field child integer[] 3 --- @field time integer 4 5 --- @alias vim.undotree.tree {[integer]: vim.undotree.tree.entry} 6 7 local M = {} 8 9 local ns = vim.api.nvim_create_namespace('nvim.undotree') 10 11 --- @param buf integer 12 --- @return vim.fn.undotree.entry[] 13 --- @return integer 14 local function get_undotree_entries(buf) 15 local undotree = vim.fn.undotree(buf) 16 local entries = undotree.entries 17 18 --Maybe: `:undo 0` and then `undotree` to get seq 0 time 19 table.insert(entries, 1, { seq = 0, time = -1 }) 20 21 return entries, undotree.seq_cur 22 end 23 24 --- @param ent vim.fn.undotree.entry[] 25 --- @param _tree vim.undotree.tree? 26 --- @param _last integer? 27 --- @return vim.undotree.tree 28 local function treefy(ent, _tree, _last) 29 local tree = _tree or {} 30 local last = _last or nil 31 32 for idx, v in ipairs(ent) do 33 local seq = v.seq 34 35 if last then 36 table.insert(tree[last].child, seq) 37 else 38 assert(idx == 1 and not _tree) 39 end 40 41 tree[seq] = { child = {}, time = v.time } 42 if v.alt then 43 assert(last) 44 treefy(v.alt, tree, last) 45 end 46 last = seq 47 end 48 49 return tree 50 end 51 52 --- @class (private) vim.undotree.graph_line 53 --- @field kind 'node'|'remove'|'branch'|'remove+branch'|'nochange_remove' 54 --- @field index integer 55 --- @field node_count integer 56 --- @field node integer|integer[] 57 --- @field index2 integer? -- for branch-index in `remove+branch` 58 59 --- @param tree vim.undotree.tree 60 --- @return vim.undotree.graph_line[] 61 local function tree_to_graph_lines(tree) 62 --- @type vim.undotree.graph_line[] 63 local graph_lines = {} 64 65 assert(tree[0], "tree doesn't have 0-th node") 66 --- @type (integer[]|integer)[] 67 local nodes = { 0 } 68 69 while #nodes > 0 do 70 local minseq = math.huge 71 --- @type integer 72 local index 73 --- @type integer 74 local node_index 75 76 for k, v in ipairs(nodes) do 77 if type(v) == 'table' then 78 for i, j in ipairs(v) do 79 if j < minseq then 80 minseq = j 81 index = k 82 node_index = i 83 end 84 end 85 elseif v < minseq then 86 assert(type(v) == 'number') 87 minseq = v 88 index = k 89 end 90 end 91 92 local node = nodes[index] 93 94 --- @param kind 'node'|'remove'|'branch'|'nochange_remove' 95 local function add_graph_line(kind) 96 table.insert(graph_lines, { kind = kind, index = index, node_count = #nodes, node = node }) 97 end 98 99 if type(node) == 'number' then 100 add_graph_line('node') 101 102 local child = tree[node].child 103 if #child == 0 then 104 if index ~= #nodes then 105 add_graph_line('remove') 106 else 107 add_graph_line('nochange_remove') 108 end 109 110 table.remove(nodes, index) 111 elseif #child == 1 then 112 nodes[index] = child[1] 113 else 114 nodes[index] = child 115 end 116 else 117 assert(type(node) == 'table') 118 119 add_graph_line('branch') 120 121 table.remove(nodes, index) 122 if #node == 2 then 123 table.insert(nodes, index, math.min(unpack(node))) 124 table.insert(nodes, index, math.max(unpack(node))) 125 elseif #node > 2 then 126 table.insert(nodes, index, node[node_index]) 127 table.insert(nodes, index, node) 128 table.remove(node, node_index) 129 end 130 end 131 end 132 133 for k, v in ipairs(graph_lines) do 134 if v.kind == 'remove' and (graph_lines[k + 1] or {}).kind == 'branch' then 135 v.kind = 'remove+branch' 136 v.index2 = graph_lines[k + 1].index 137 table.remove(graph_lines, k + 1) 138 end 139 end 140 141 return graph_lines 142 end 143 144 --- @param time integer 145 --- @return string 146 local function undo_fmt_time(time) 147 if time == -1 then 148 return 'origin' 149 end 150 151 local diff = os.time() - time 152 153 if diff >= 100 then 154 if diff < (60 * 60 * 12) then 155 return os.date('%H:%M:%S', time) --[[@as string]] 156 else 157 return os.date('%Y/%m/%d %H:%M:%S', time) --[[@as string]] 158 end 159 else 160 return ('%d second%s ago'):format(diff, diff == 1 and '' or 's') 161 end 162 end 163 164 --- @param tree vim.undotree.tree 165 --- @param graph_lines vim.undotree.graph_line[] 166 --- @param buf integer 167 --- @param meta {[integer]:integer} 168 --- @param find_seq? integer 169 --- @return integer? 170 local function buf_apply_graph_lines(tree, graph_lines, buf, meta, find_seq) 171 -- As in io-buffer, not vim-buffer 172 local line_buffer = {} 173 local extmark_buffer = {} 174 175 --- @type integer? 176 local found_seq 177 178 for k, v in ipairs(graph_lines) do 179 local is_last = k == #graph_lines 180 181 --- @type string? 182 local line 183 if v.kind == 'node' then 184 line = ('| '):rep(v.index - 1) 185 .. '*' 186 .. (' |'):rep(v.node_count - v.index) 187 .. ' ' 188 .. v.node 189 .. ' (' 190 .. undo_fmt_time(tree[v.node].time) 191 .. ')' 192 elseif v.kind == 'remove' then 193 line = ('| '):rep(v.index - 1) .. (' /'):rep(v.node_count - v.index) 194 elseif v.kind == 'branch' then 195 line = ('| '):rep(v.index - 1) .. '|\\' .. (' \\'):rep(v.node_count - v.index) 196 elseif v.kind == 'remove+branch' then 197 if v.index2 < v.index then 198 line = ('| '):rep(v.index2 - 1) 199 .. '|\\' 200 .. (' \\'):rep(v.index - v.index2 - 1) 201 .. ' ' 202 .. (' |'):rep(v.node_count - v.index) 203 else 204 line = ('| '):rep(v.index - 1) 205 .. (' /'):rep(v.index2 - v.index) 206 .. ' /|' 207 .. (' |'):rep(v.node_count - v.index2 - 1) 208 end 209 elseif v.kind == 'nochange_remove' then 210 line = nil 211 else 212 error 'unreachable' 213 end 214 215 if v.kind == 'node' then 216 table.insert(line_buffer, line) 217 table.insert(meta, v.node) 218 219 if v.node == find_seq then 220 found_seq = #meta 221 end 222 elseif line then 223 table.insert(extmark_buffer, { { line, 'Normal' } }) 224 end 225 226 if next(extmark_buffer) and (v.kind == 'node' or is_last) then 227 local row = vim.api.nvim_buf_line_count(buf) 228 vim.api.nvim_buf_set_extmark(buf, ns, row - 1, 0, { virt_lines = extmark_buffer }) 229 extmark_buffer = {} 230 end 231 232 if next(line_buffer) and (v.kind ~= 'node' or is_last) then 233 vim.api.nvim_buf_set_lines(buf, -1, -1, true, line_buffer) 234 235 if #line_buffer > 3 then 236 local end_ = vim.api.nvim_buf_line_count(buf) - 1 237 local start = end_ - #line_buffer + 3 238 vim.api.nvim_buf_call(buf, function() 239 local w = vim.b[buf].nvim_is_undotree 240 if vim.api.nvim_win_is_valid(w) and vim.wo[w].foldmethod == 'manual' then 241 vim.cmd.fold { range = { start, end_ } } 242 end 243 end) 244 end 245 246 line_buffer = {} 247 end 248 end 249 250 vim.api.nvim_buf_set_lines(buf, 0, 1, true, {}) 251 252 return found_seq 253 end 254 255 ---@param inbuf integer 256 ---@param outbuf integer 257 ---@return {[integer]:integer} 258 local function draw(inbuf, outbuf) 259 local entries, curseq = get_undotree_entries(inbuf) 260 local tree = treefy(entries) 261 local graph_lines = tree_to_graph_lines(tree) 262 263 local meta = {} 264 vim.bo[outbuf].modifiable = true 265 vim.api.nvim_buf_set_lines(outbuf, 0, -1, true, {}) 266 vim.api.nvim_buf_clear_namespace(outbuf, ns, 0, -1) 267 local curseq_line = buf_apply_graph_lines(tree, graph_lines, outbuf, meta, curseq) 268 vim.bo[outbuf].modifiable = false 269 270 vim.schedule(function() 271 if vim.api.nvim_win_is_valid(vim.b[outbuf].nvim_is_undotree) then 272 vim.api.nvim_win_set_cursor(vim.b[outbuf].nvim_is_undotree, { curseq_line, 0 }) 273 end 274 end) 275 276 return meta 277 end 278 279 --- @class vim.undotree.opts 280 --- @inlinedoc 281 --- 282 --- Buffer to draw the tree into. If omitted, a new buffer is created. 283 --- @field bufnr integer? 284 --- 285 --- Window id to display the tree buffer in. If omitted, a new window is 286 --- created with {command}. 287 --- @field winid integer? 288 --- 289 --- Vimscript command to create the window. Default value is "30vnew". 290 --- Only used when {winid} is nil. 291 --- @field command string? 292 --- 293 --- Title of the window. If a function, it accepts the buffer number of the 294 --- source buffer as its only argument and should return a string. 295 --- @field title (string|fun(bufnr:integer):string|nil)? 296 297 --- Open a window that displays a textual representation of the [undo-tree]. 298 --- 299 --- While in the window, moving the cursor changes the undo. 300 --- 301 --- Closes the window if it is already open 302 --- 303 --- Load the plugin with this command: 304 --- ``` 305 --- packadd nvim.undotree 306 --- ``` 307 --- 308 --- Can also be shown with `:Undotree`. [:Undotree]() 309 --- 310 --- @param opts vim.undotree.opts? 311 --- @return boolean? Returns true if the window was already open, nil otherwise 312 function M.open(opts) 313 -- The following lines of code was copied from 314 -- `vim.treesitter.dev.inspect_tree` and then modified to fit 315 316 vim.validate('opts', opts, 'table', true) 317 318 opts = opts or {} 319 320 local buf = vim.api.nvim_get_current_buf() 321 322 if vim.b[buf].nvim_undotree then 323 local w = vim.b[buf].nvim_undotree 324 if vim.api.nvim_win_is_valid(w) then 325 vim.api.nvim_win_close(w, true) 326 return true 327 end 328 elseif vim.b[buf].nvim_is_undotree then 329 local w = vim.b[buf].nvim_is_undotree 330 if vim.api.nvim_win_is_valid(w) then 331 vim.api.nvim_win_close(w, true) 332 return true 333 end 334 end 335 336 local w = opts.winid 337 if not w then 338 vim.cmd(opts.command or '30vnew') 339 w = vim.api.nvim_get_current_win() 340 end 341 342 local b = opts.bufnr 343 if b then 344 vim.api.nvim_win_set_buf(w, b) 345 else 346 b = vim.api.nvim_win_get_buf(w) 347 end 348 349 vim.b[buf].nvim_undotree = w 350 vim.b[b].nvim_is_undotree = w 351 352 local title --- @type string? 353 local opts_title = opts.title 354 if not opts_title then 355 local bufname = vim.api.nvim_buf_get_name(buf) 356 title = string.format('Undo tree for %s', vim.fn.fnamemodify(bufname, ':.')) 357 elseif type(opts_title) == 'function' then 358 title = opts_title(buf) 359 elseif type(opts_title) == 'string' then 360 title = opts_title 361 end 362 363 assert(type(title) == 'string', 'Window title must be a string') 364 vim.api.nvim_buf_set_name(b, title) 365 366 vim.wo[w][0].scrolloff = 5 367 vim.wo[w][0].wrap = false 368 vim.wo[w][0].foldmethod = 'manual' 369 vim.wo[w][0].foldenable = true 370 vim.wo[w][0].cursorline = true 371 vim.bo[b].buflisted = false 372 vim.bo[b].buftype = 'nofile' 373 vim.bo[b].bufhidden = 'wipe' 374 vim.bo[b].swapfile = false 375 376 local meta = draw(buf, b) 377 378 vim.api.nvim_win_set_cursor(w, { vim.api.nvim_buf_line_count(b), 0 }) 379 380 local group = vim.api.nvim_create_augroup('nvim.undotree', {}) 381 382 vim.api.nvim_win_call(w, function() 383 vim.cmd.syntax('region Comment start="(" end=")"') 384 end) 385 386 vim.api.nvim_create_autocmd('CursorMoved', { 387 group = group, 388 buffer = b, 389 callback = function() 390 local row = vim.fn.line('.') 391 vim.api.nvim_buf_call(buf, function() 392 vim.cmd.undo { meta[row], mods = { silent = true } } 393 end) 394 end, 395 }) 396 397 vim.api.nvim_create_autocmd({ 'TextChanged', 'InsertLeave' }, { 398 group = group, 399 buffer = buf, 400 callback = function() 401 if not vim.api.nvim_buf_is_valid(b) then 402 return true 403 end 404 405 meta = draw(buf, b) 406 407 if vim.api.nvim_win_is_valid(w) then 408 vim.wo[w][0].foldlevel = 99 409 end 410 end, 411 }) 412 413 vim.bo[b].filetype = 'nvim-undotree' 414 end 415 416 return M