restructure tree visit algorithm

lua/u/renderer.lua

@@ -1,15 +1,25 @@
+---@diagnostic disable: redefined-local
 function _G.URendererOpFuncSwallow() end
 
-local ENABLE_LOG = false
+--- Comma-separated list of regex patterns for debug logging. Examples:
+--- "render,extmark" or ".*" or "get_tags_at,_on_text_changed" or
+--- "^_on_.*,reconcile$"
+local DEBUG = vim.env.DEBUG or ''
 
-local function log(...)
-  if not ENABLE_LOG then return end
+--- @param name string
+local function log(name, ...)
+  if DEBUG == '' then return end
 
+  for pattern in DEBUG:gmatch '[^,]+' do
+    if name:match(vim.trim(pattern)) then
       local f = assert(
         io.open(vim.fs.joinpath((vim.fn.stdpath 'log') --[[@as string]], 'u.renderer.log'), 'a+')
       )
-  f:write(os.date() .. '\t' .. vim.inspect { ... } .. '\n')
+      f:write(os.date() .. '\t' .. vim.inspect { name, ... } .. '\n')
       f:close()
+      return
+    end
+  end
 end
 
 local M = {}
@@ -22,12 +32,33 @@ local H = {}
 --- @generic TProps
 --- @generic TState
 --- @class u.renderer.FnComponentContext<TProps, TState>
---- @field props TProps
---- @field state u.Signal<TState>
---- @field children u.renderer.Tree
 --- @field phase 'mount'|'update'|'unmount'
---- @field private unsubscribe? fun(): any
---- @field private prev_tree u.renderer.Tree
+--- @field props TProps
+--- @field state? TState
+--- @field children u.renderer.Tree
+--- @field private on_change? fun(): any
+--- @field private prev_rendered_children? u.renderer.Tree
+local FnComponentContext = {}
+M.FnComponentContext = FnComponentContext
+FnComponentContext.__index = FnComponentContext
+
+--- @param props TProps
+--- @param state? TState
+--- @param children u.renderer.Tree
+function FnComponentContext.new(props, state, children)
+  return setmetatable({
+    phase = 'mount',
+    props = props,
+    state = state,
+    children = children,
+  }, FnComponentContext)
+end
+
+--- @param new_state TState
+function FnComponentContext:update(new_state)
+  self.state = new_state
+  if self.on_change then vim.schedule(function() self.on_change() end) end
+end
 
 --- @alias u.renderer.FnComponent<TProps, TState> fun(ctx: u.renderer.FnComponentContext<TProps, TState>): u.renderer.Tree
 
@@ -36,6 +67,7 @@ local H = {}
 --- @field name string | u.renderer.FnComponent<any, any>
 --- @field attributes u.renderer.TagAttributes
 --- @field children u.renderer.Tree
+--- @field private ctx? u.renderer.FnComponentContext
 
 --- @alias u.renderer.Node nil | boolean | string | u.renderer.Tag
 --- @alias u.renderer.Tree u.renderer.Node | u.renderer.Node[]
@@ -70,7 +102,7 @@ M.h = setmetatable({}, {
 --- @field ns integer
 --- @field changedtick integer
 --- @field text_content { old: { lines: string[], extmarks: u.renderer.RendererExtmark[] }, curr: { lines: string[], extmarks: u.renderer.RendererExtmark[] } }
---- @field component_tree { old: u.renderer.Tree, ctx_by_node: table<u.renderer.Tree, u.renderer.FnComponentContext> }
+--- @field component_tree { old: u.renderer.Tree  }
 local Renderer = {}
 Renderer.__index = Renderer
 M.Renderer = Renderer
@@ -391,9 +423,9 @@ function Renderer:_expr_map_callback(mode, lhs) -- {{{
   -- find the tag with the smallest intersection that contains the cursor:
   local pos0 = vim.api.nvim_win_get_cursor(0)
   pos0[1] = pos0[1] - 1 -- make it actually 0-based
-  log('_expr_map_callback: pos0:', pos0)
+  log('_expr_map_callback', 'pos0:', pos0)
   local pos_infos = self:get_tags_at(pos0)
-  log('_expr_map_callback: pos_infos:', pos_infos)
+  log('_expr_map_callback', 'pos_infos:', pos_infos)
 
   if #pos_infos == 0 then return lhs end
 
@@ -443,7 +475,7 @@ function Renderer:_on_text_changed() -- {{{
       --- @type integer, integer, vim.api.keyset.extmark_details
       local start_row0, start_col0, details = unpack(extmark)
       local end_row0, end_col0 = details.end_row, details.end_col
-      log('_on_text_changed: fetched current extmark for pos_info', {
+      log('_on_text_changed', 'fetched current extmark for pos_info', {
         pos_info = pos_info,
         curr_extmark = {
           start_row0 = start_row0,
@@ -472,7 +504,7 @@ function Renderer:_on_text_changed() -- {{{
           or ''
         end_col0 = last_line:len()
       end
-      log('_on_text_changed: after position correction', {
+      log('_on_text_changed', 'after position correction', {
         curr_extmark = {
           start_row0 = start_row0,
           start_col0 = start_col0,
@@ -490,7 +522,7 @@ function Renderer:_on_text_changed() -- {{{
       local pos2 = { self.bufnr, end_row0 + 1, end_col0 }
       local ok, lines = pcall(vim.fn.getregion, pos1, pos2, { type = 'v' })
       if not ok then
-        log('_on_text_changed: getregion: invalid-pos ', {
+        log('_on_text_changed', 'getregion: invalid-pos ', {
           { pos1, pos2 },
         })
         vim.api.nvim_echo({
@@ -601,7 +633,8 @@ function Renderer:get_tags_at(pos0, mode) -- {{{
     { details = true, overlap = true }
   )
   log(
-    'get_tags_at: context:',
+    'get_tags_at',
+    'context:',
     { pos0 = pos0, mode = mode, raw_overlapping_extmarks = raw_overlapping_extmarks }
   )
 
@@ -719,43 +752,42 @@ end -- }}}
 
 --- @param tree u.renderer.Tree
 function Renderer:mount(tree) -- {{{
-  local tracker = require 'u.tracker'
-
   local function render()
     local H2 = {}
 
     --- @param tree u.renderer.Tree
     function H2.unmount(tree)
+      log('Render.mount.umount', { tree = tree })
       --- @param tree u.renderer.Tree
-      local function unmount_one(tree)
+      local function visit(tree)
         Renderer.tree_match(tree, {
           array = function(tags)
+            log('Render.mount.umount.array', { tags = tags })
             for _, tag in ipairs(tags) do
-              unmount_one(tag)
+              visit(tag)
             end
           end,
-          tag = function(tag) unmount_one(tag.children) end,
+          tag = function(tag)
+            log('Render.mount.umount.tag', { tag = tag })
+            visit(tag.children)
+          end,
           component = function(Component, tag)
+            log('Render.mount.umount.component', { Component = Component, tag = tag })
             --- @type u.renderer.FnComponentContext
-            local ctx = assert(
-              self.component_tree.ctx_by_node[tree] --[[@as any]],
-              'could not find context for node'
-            )
+            local ctx = assert(tag.ctx, 'could not find context for node')
 
             -- depth-first:
-            unmount_one(ctx.children)
+            visit(ctx.children)
 
             -- now unmount current:
             ctx.phase = 'unmount'
             Component(ctx)
-            local unsubscribe = assert(ctx.unsubscribe, 'unmount: unsubscribe is nil')
-            unsubscribe()
-            self.component_tree.ctx_by_node[tree] = nil
+            ctx.on_change = nil
           end,
         })
       end
 
-      unmount_one(tree)
+      visit(tree)
     end
 
     --- @param old_tree u.renderer.Tree
@@ -764,6 +796,12 @@ function Renderer:mount(tree) -- {{{
     function H2.visit_tree(old_tree, new_tree)
       local old_tree_kind = Renderer.tree_kind(old_tree)
       local new_tree_kind = Renderer.tree_kind(new_tree)
+      log('Renderer.mount.visit_tree', {
+        old_tree_kind = old_tree_kind,
+        old_tree = old_tree,
+        new_tree_kind = new_tree_kind,
+        new_tree = new_tree,
+      })
 
       local new_tree_rendered = Renderer.tree_match(new_tree, {
         string = function(s) return s end,
@@ -774,6 +812,7 @@ function Renderer:mount(tree) -- {{{
           return H2.visit_array(old_tree --[[@as any]], new_arr)
         end,
         tag = function(new_tag)
+          log('Renderer.mount.visit_tree.tag', { new_tag = new_tag })
           local old_children = old_tree_kind == new_tree_kind and old_tree.children or nil
           return M.h(
             new_tag.name,
@@ -783,47 +822,34 @@ function Renderer:mount(tree) -- {{{
         end,
 
         component = function(NewC, new_tag)
-          local old_ctx = old_tree_kind == new_tree_kind
-              and self.component_tree.ctx_by_node[old_tree]
-            or nil
-          if old_tree_kind == new_tree_kind then self.component_tree.ctx_by_node[old_tree] = nil end
+          --- @type { tag: u.renderer.Tag, ctx?: u.renderer.FnComponentContext } | nil
+          local old_component_info = Renderer.tree_match(old_tree, { component = function(_, t) return { tag = t, ctx = t.ctx } end })
+          local old_tag = old_component_info and old_component_info.tag or nil
+          local ctx = old_component_info and old_component_info.ctx or nil
 
-          --- @param ctx u.renderer.FnComponentContext
-          --- @param rendered_component u.renderer.Tree
-          local function update_ctx_and_recurse(ctx, rendered_component)
-            self.component_tree.ctx_by_node[new_tree] = ctx
-            -- The rendered component could have other components in its tree:
-            -- recurse and simplify:
-            local new_tree_simplified =
-              M.h('text', ctx.props, H2.visit_tree(ctx.prev_tree, rendered_component))
-            ctx.prev_tree = rendered_component
-            return new_tree_simplified
-          end
+          -- TODO: if the left doesn't match, do we unmount here, or leave that for the later code?
 
-          if old_ctx then
-            old_ctx.phase = 'update'
-            old_ctx.props = new_tag.attributes
-            old_ctx.children = new_tag.children
-            return update_ctx_and_recurse(old_ctx, NewC(old_ctx))
-          else
-            local state = tracker.create_signal(nil)
+          log('Renderer.mount.visit_tree.component', {
+            NewC = NewC,
+            new_tag = new_tag,
+            old_tag = old_tag,
+          })
 
+          if not ctx then
             --- @type u.renderer.FnComponentContext
-            local ctx = {
-              phase = 'mount',
-              props = new_tag.attributes,
-              children = new_tag.children,
-              state = state --[[@as any]],
-              prev_tree = nil,
-            }
-            local rendered = NewC(ctx)
-            -- This is the main difference between this else, and the above if:
-            -- we subscribe _after_ first render. This makes it convenient for
-            -- first-renders to initialize the state without firing a
-            -- re-render:
-            ctx.unsubscribe = state:subscribe(render)
-            return update_ctx_and_recurse(ctx, rendered)
+            ctx = FnComponentContext.new(new_tag.attributes, nil, new_tag.children)
+          else
+            ctx.phase = 'update'
           end
+          ctx.props = new_tag.attributes
+          ctx.children = new_tag.children
+          ctx.on_change = render
+
+          new_tag.ctx = ctx
+          local NewC_rendered_children = NewC(ctx)
+          local result = H2.visit_tree(ctx.prev_rendered_children, NewC_rendered_children)
+          ctx.prev_rendered_children = NewC_rendered_children
+          return result
         end,
       })
 
@@ -836,7 +862,7 @@ function Renderer:mount(tree) -- {{{
     --- @param new_arr u.renderer.Node[]
     --- @return u.renderer.Node[]
     function H2.visit_array(old_arr, new_arr)
-      --- @type table<any, u.renderer.Node>
+      --- @type table<any, u.renderer.Node | nil>
       local old_by_key = {}
 
       --- @param tree u.renderer.Tree
@@ -871,6 +897,12 @@ function Renderer:mount(tree) -- {{{
         table.insert(resulting_nodes, H2.visit_tree(old, new))
       end
 
+      log('Renderer.mount.visit_array', {
+        old_arr = old_arr,
+        new_arr = new_arr,
+        resulting_nodes = resulting_nodes,
+      })
+
       -- Unmount whatever is left in the "old" map:
       for _, n in pairs(old_by_key) do
         H2.unmount(n)