wip

Author: MunifTanjim

lua/nui/tree/init.lua

@@ -1,7 +1,6 @@
 local Object = require("nui.object")
 local _ = require("nui.utils")._
 local defaults = require("nui.utils").defaults
-local is_type = require("nui.utils").is_type
 local tree_util = require("nui.tree.util")
 
 -- returns id of the first window that contains the buffer
@@ -11,11 +10,12 @@ local function get_winid(bufnr)
   return vim.fn.win_findbuf(bufnr)[1]
 end
 
+---@param get_tree fun(): NuiTree
 ---@param nodes NuiTree.Node[]
 ---@param parent_node? NuiTree.Node
 ---@param get_node_id nui_tree_get_node_id
 ---@return { by_id: table<string, NuiTree.Node>, root_ids: string[] }
-local function initialize_nodes(nodes, parent_node, get_node_id)
+local function initialize_nodes(get_tree, nodes, parent_node, get_node_id)
   local start_depth = parent_node and parent_node:get_depth() + 1 or 1
 
   ---@type table<string, NuiTree.Node>
@@ -26,7 +26,9 @@ local function initialize_nodes(nodes, parent_node, get_node_id)
   ---@param node NuiTree.Node
   ---@param depth number
   local function initialize(node, depth)
+    node._tree = get_tree
     node._depth = depth
+    node._height = 0
     node._id = get_node_id(node)
     node._initialized = true
 
@@ -75,6 +77,8 @@ end
 ---@field _depth integer
 ---@field _parent_id? string
 ---@field _child_ids? string[]
+---@field _next string
+---@field _prev string
 ---@field __children? NuiTree.Node[]
 ---@field [string] any
 local TreeNode = {
@@ -147,6 +151,9 @@ end
 ---@field node_id_by_linenr table<integer, string>
 ---@field prepare_node nui_tree_prepare_node
 ---@field win_options table<string, any> # deprecated
+---@field head? string
+---@field tail? string
+---@field get_tree fun(): NuiTree
 
 ---@class nui_tree_options
 ---@field bufnr integer
@@ -212,6 +219,10 @@ function Tree:init(options)
     linenr = {},
   }
 
+  self._.get_tree = function()
+    return self
+  end
+
   _.set_buf_options(self.bufnr, self._.buf_options)
 
   ---@deprecated
@@ -253,14 +264,48 @@ end
 ---@return nil|integer linenr
 ---@return nil|integer linenr
 function Tree:get_node(node_id_or_linenr)
-  if is_type("string", node_id_or_linenr) then
-    return self.nodes.by_id[node_id_or_linenr], unpack(self._.linenr_by_node_id[node_id_or_linenr] or {})
+  local by_id = self.nodes.by_id
+
+  if type(node_id_or_linenr) == "string" then
+    local node = by_id[node_id_or_linenr]
+    if not node then
+      return
+    end
+
+    local linenr = 1
+    local head = by_id[self._head]
+    while head and head ~= node do
+      head = by_id[head._next]
+      linenr = linenr + 1
+    end
+
+    return node, linenr, linenr + node._height - 1
+
+    -- return node, unpack(self._.linenr_by_node_id[node_id_or_linenr] or {})
   end
 
   local winid = get_winid(self.bufnr)
   local linenr = node_id_or_linenr or vim.api.nvim_win_get_cursor(winid)[1]
-  local node_id = self._.node_id_by_linenr[linenr]
-  return self.nodes.by_id[node_id], unpack(self._.linenr_by_node_id[node_id] or {})
+
+  local node, linenr_left = by_id[self._head], linenr - 1
+  local node_height_below = node and node._height
+  while node and linenr_left > 0 do
+    linenr_left = linenr_left - 1
+    node_height_below = node_height_below - 1
+    if node_height_below == 0 then
+      node = by_id[node._next]
+      node_height_below = node and node._height
+    end
+  end
+
+  if not node then
+    return
+  end
+
+  return node, linenr + node_height_below - node._height, linenr + node_height_below - 1
+
+  -- local node_id = self._.node_id_by_linenr[linenr]
+  -- return self.nodes.by_id[node_id], unpack(self._.linenr_by_node_id[node_id] or {})
 end
 
 ---@param parent_id? string parent node's id
@@ -285,7 +330,7 @@ end
 ---@param nodes NuiTree.Node[]
 ---@param parent_node? NuiTree.Node
 function Tree:_add_nodes(nodes, parent_node)
-  local new_nodes = initialize_nodes(nodes, parent_node, self._.get_node_id)
+  local new_nodes = initialize_nodes(self._.get_tree, nodes, parent_node, self._.get_node_id)
 
   self.nodes.by_id = vim.tbl_extend("force", self.nodes.by_id, new_nodes.by_id)
 
@@ -307,8 +352,8 @@ end
 ---@param nodes NuiTree.Node[]
 ---@param parent_id? string parent node's id
 function Tree:set_nodes(nodes, parent_id)
-  self._.node_id_by_linenr = {}
-  self._.linenr_by_node_id = {}
+  -- self._.node_id_by_linenr = {}
+  -- self._.linenr_by_node_id = {}
 
   if not parent_id then
     self.nodes = { by_id = {}, root_ids = {} }
@@ -387,11 +432,13 @@ function Tree:_prepare_content(linenr_start)
   local tree_linenr = 0
   local lines = { len = tree_linenr }
 
-  local node_id_by_linenr = {}
-  internal.node_id_by_linenr = node_id_by_linenr
+  -- local node_id_by_linenr = {}
+  -- internal.node_id_by_linenr = node_id_by_linenr
+
+  -- local linenr_by_node_id = {}
+  -- internal.linenr_by_node_id = linenr_by_node_id
 
-  local linenr_by_node_id = {}
-  internal.linenr_by_node_id = linenr_by_node_id
+  local prev_node = nil
 
   local function prepare(node_id, parent_node)
     local node = by_id[node_id]
@@ -400,30 +447,42 @@ function Tree:_prepare_content(linenr_start)
     end
 
     local node_lines = internal.prepare_node(node, parent_node)
+    local node_lines_len = 0
     if node_lines then
       if type(node_lines) ~= "table" or node_lines.content then
         list_wrapper[1] = node_lines
         node_lines = list_wrapper
       end
       ---@cast node_lines -string, -NuiLine
 
-      local node_linenr = linenr_by_node_id[node_id] or {}
-      for node_line_idx = 1, #node_lines do
+      -- local node_linenr = linenr_by_node_id[node_id] or {}
+      node_lines_len = #node_lines
+      for node_line_idx = 1, node_lines_len do
         local node_line = node_lines[node_line_idx]
 
         tree_linenr = tree_linenr + 1
-        local buffer_linenr = tree_linenr + linenr_start - 1
+        -- local buffer_linenr = tree_linenr + linenr_start - 1
 
         lines[tree_linenr] = node_line
 
-        node_id_by_linenr[buffer_linenr] = node_id
+        -- node_id_by_linenr[buffer_linenr] = node_id
+
+        -- if node_line_idx == 1 then
+        --   node_linenr[1] = buffer_linenr
+        -- end
+        -- node_linenr[2] = buffer_linenr
+      end
+      -- linenr_by_node_id[node_id] = node_linenr
+    end
+    node._height = node_lines_len
 
-        if node_line_idx == 1 then
-          node_linenr[1] = buffer_linenr
-        end
-        node_linenr[2] = buffer_linenr
+    if node._height > 0 then
+      node._prev = prev_node
+      if prev_node then
+        node._prev = node._prev:get_id()
+        prev_node._next = node:get_id()
       end
-      linenr_by_node_id[node_id] = node_linenr
+      prev_node = node
     end
 
     local child_ids = node._child_ids
@@ -439,6 +498,9 @@ function Tree:_prepare_content(linenr_start)
     prepare(root_ids[node_id_idx])
   end
 
+  self._head = root_ids[1]
+  self._tail = prev_node
+
   lines.len = tree_linenr
 
   return lines