Merge pull request #91 from VEZY/fix-#90

Fix: issue when updating node id or merging trees

Author: VEZY

src/compute_MTG/indexing.jl

@@ -79,10 +79,23 @@ end
     attrs = node_attributes(node)
     store = attrs.ref.store
     store === nothing && return nothing
+    if store !== plan.store
+        throw(ArgumentError(
+            "Incoherent columnar query plan for key `$(key)` on node id $(node_id(node)): " *
+            "the node belongs to a different attribute store than the query root. " *
+            "This usually happens after attaching subtrees from independent MTGs. " *
+            "Rebuild a unified store with `MultiScaleTreeGraph.columnarize!(get_root(node))`."
+        ))
+    end
     nodeid = node_id(node)
     nodeid > length(store.node_bucket) && return nothing
     bid = store.node_bucket[nodeid]
     bid == 0 && return nothing
+    bid > length(plan.col_idx_by_bucket) && throw(ArgumentError(
+        "Incoherent columnar query plan for key `$(key)` on node id $(nodeid): " *
+        "bucket id $(bid) is outside query-plan bounds $(length(plan.col_idx_by_bucket)). " *
+        "Rebuild a unified store with `MultiScaleTreeGraph.columnarize!(get_root(node))`."
+    ))
     col_idx = plan.col_idx_by_bucket[bid]
     col_idx == 0 && return nothing
     row = store.node_row[nodeid]

src/compute_MTG/node_funs.jl

@@ -77,14 +77,33 @@ function addchild!(p::Node, MTG::M) where {M<:AbstractNodeMTG}
     return child
 end
 
+@inline function _columnar_store_or_nothing(node::Node)
+    attrs = node_attributes(node)
+    attrs isa ColumnarAttrs || return nothing
+    return _store_for_node_attrs(attrs)
+end
+
+@inline function _maybe_recolumnarize_after_attach!(p::Node, child::Node, child_was_root::Bool)
+    child_was_root || return nothing
+    p_store = _columnar_store_or_nothing(p)
+    c_store = _columnar_store_or_nothing(child)
+    if p_store !== nothing && c_store !== nothing && p_store !== c_store
+        columnarize!(get_root(p))
+    end
+    return nothing
+end
+
 function addchild!(p::Node{N,A}, child::Node; force=false) where {N<:AbstractNodeMTG,A}
-    if parent(child) === nothing || force == true
+    child_was_root = parent(child) === nothing
+
+    if child_was_root || force == true
         reparent!(child, p)
     elseif parent(child) != p && force == false
         error("The node already has a parent. Hint: use `force=true` if needed.")
     end
 
     push!(children(p), child)
+    _maybe_recolumnarize_after_attach!(p, child, child_was_root)
 
     return child
 end

src/types/Attributes.jl

@@ -45,6 +45,7 @@ end
 
 struct ColumnarQueryPlan
     key::Symbol
+    store::MTGAttributeStore
     col_idx_by_bucket::Vector{Int}
 end
 
@@ -110,7 +111,25 @@ function _rebuild_subtree_index!(store::MTGAttributeStore, root)
         pos = stack_pos[end]
 
         if pos == 0
+            current_attrs = node_attributes(current)
+            current_attrs isa ColumnarAttrs || throw(ArgumentError(
+                "Subtree index rebuild requires a columnar attribute backend."
+            ))
+            current_store = _store_for_node_attrs(current_attrs)
+            if current_store !== store
+                throw(ArgumentError(
+                    "Incoherent MTG columnar stores detected while building descendants index: " *
+                    "node id $(node_id(current)) belongs to a different attribute store than the root. " *
+                    "Rebuild a unified store with `MultiScaleTreeGraph.columnarize!(get_root(root))`."
+                ))
+            end
             nid = node_id(current)
+            if nid > nmax || store.node_bucket[nid] == 0
+                throw(ArgumentError(
+                    "Incoherent MTG columnar indexing state for node id $(nid). " *
+                    "Rebuild a unified store with `MultiScaleTreeGraph.columnarize!(get_root(root))`."
+                ))
+            end
             t += 1
             tin[nid] = t
             push!(dfs_order, nid)
@@ -398,7 +417,7 @@ function build_columnar_query_plan(node, key::Symbol)
     @inbounds for i in eachindex(store.buckets)
         col_idx_by_bucket[i] = get(store.buckets[i].col_index, key, 0)
     end
-    ColumnarQueryPlan(key, col_idx_by_bucket)
+    ColumnarQueryPlan(key, store, col_idx_by_bucket)
 end
 
 @inline function _symbol_bucket_ids(store::MTGAttributeStore, symbol_filter)

test/test-descendants.jl

@@ -65,6 +65,25 @@
       @test descendants!(out_nodes, get_node(mtg, 6), self=true) == [get_node(mtg, 6), get_node(mtg, 7)]
 end
 
+@testset "descendants clear error on mixed columnar stores" begin
+      mtg_a = read_mtg("files/simple_plant.mtg")
+      mtg_b = read_mtg("files/simple_plant.mtg")
+      # Bypass addchild! on purpose to build an incoherent tree (mixed stores).
+      reparent!(mtg_b, mtg_a)
+      push!(children(mtg_a), mtg_b)
+
+      err = try
+            descendants(mtg_a, :Width)
+            nothing
+      catch e
+            e
+      end
+
+      @test err isa ArgumentError
+      @test occursin("different attribute store", sprint(showerror, err))
+      @test occursin("columnarize!", sprint(showerror, err))
+end
+
 # using BenchmarkTools
 # @benchmark descendants($mtg, :Width) # 876 ns
 # @benchmark descendants!($mtg, :Width) # 5.6 μs

test/test-nodes.jl

@@ -126,3 +126,18 @@ end
     mtg = read_mtg(file, MutableNodeMTG)
     VERSION >= v"1.7" && @test_throws "The parent node has an MTG encoding of type `MutableNodeMTG`, but the MTG encoding you provide is of type `NodeMTG`, please make sure they are the same." Node(mtg, NodeMTG(:/, :Branch, 1, 2))
 end
+
+@testset "addchild! re-columnarizes when attaching a root subtree" begin
+    mtg_a = read_mtg(file)
+    mtg_b = read_mtg(file)
+    addchild!(mtg_a, mtg_b; force=true)
+    @test_nowarn descendants(mtg_a, :Width)
+end
+
+@testset "new child node attributes are in the columnar store" begin
+    mtg = read_mtg(file)
+    leaf = addchild!(mtg, MutableNodeMTG(:+, :Leaf, 999, 2), Dict{Symbol,Any}(:my_attr => 42))
+    @test leaf[:my_attr] == 42
+    vals = descendants(mtg, :my_attr; ignore_nothing=true)
+    @test vals == [42]
+end