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
@@ -15,7 +14,7 @@ end
 ---@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(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 +25,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 = tree
     node._depth = depth
+    node._height = 0
     node._id = get_node_id(node)
     node._initialized = true
 
@@ -75,6 +76,8 @@ end
 ---@field _depth integer
 ---@field _parent_id? string
 ---@field _child_ids? string[]
+---@field _next NuiTree.Node
+---@field _prev NuiTree.Node
 ---@field __children? NuiTree.Node[]
 ---@field [string] any
 local TreeNode = {
@@ -147,6 +150,8 @@ 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? NuiTree.Node
+---@field tail? NuiTree.Node
 
 ---@class nui_tree_options
 ---@field bufnr integer
@@ -253,14 +258,46 @@ 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 {})
+  if type(node_id_or_linenr) == "string" then
+    local node = self.nodes.by_id[node_id_or_linenr]
+    if not node then
+      return
+    end
+
+    local linenr = 1
+    local head = self._head
+    while head and head ~= node do
+      head = 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 = 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 = 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 +322,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, 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 +344,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 +424,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 +439,41 @@ 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
+        -- 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
+      -- linenr_by_node_id[node_id] = node_linenr
+    end
+    node._height = node_lines_len
+
+    if node._height > 0 then
+      node._prev = prev_node
+      if prev_node then
+        prev_node._next = node
+      end
+      prev_node = node
     end
 
     local child_ids = node._child_ids
@@ -439,6 +489,9 @@ function Tree:_prepare_content(linenr_start)
     prepare(root_ids[node_id_idx])
   end
 
+  self._head = by_id[root_ids[1]]
+  self._tail = prev_node
+
   lines.len = tree_linenr
 
   return lines