feat: Add in-place TTree branch addition and row extension

src/uproot/writing/_cascadetree.py

@@ -96,6 +96,9 @@ def __init__(
         self._field_name = field_name
         self._basket_capacity = initial_basket_capacity
         self._resize_factor = resize_factor
+        self._has_unsupported_branches = False
+        self._has_divergent_baskets = False
+        self._needs_relocation_before_extend = False
 
         if isinstance(branch_types, dict):
             branch_types_items = branch_types.items()
@@ -446,7 +449,55 @@ def num_entries(self):
     def num_baskets(self):
         return self._num_baskets
 
+    def _relocate(self, file, sink):
+        """
+        Rewrites this tree's whole metadata blob from scratch via write_anew(),
+        relocating it (freeing its old space, allocating new space), and moves
+        the WritableFile._trees cache entry to match. Shared by the two
+        situations that need a full rewrite rather than an incremental
+        write_updates() patch: basket-capacity growth, and (see
+        _needs_relocation_before_extend) establishing Uproot's own canonical
+        layout for a tree write_updates() hasn't verified it can trust yet.
+        """
+        # seek_location is the tree's actual physical on-disk position (it's
+        # what every sink.read(key.seek_location, ...) call in
+        # _load_existing_ttree uses), and what file._move_tree()/
+        # WritableFile._trees are keyed by everywhere else in the codebase.
+        # location coincides with it for any tree Uproot itself has written
+        # (all our own tests, hence never catching this), but can diverge for
+        # a preexisting (e.g. ROOT-written) tree's Key object -- silently
+        # freeing/keying the wrong byte range and cache entry.
+        oldloc = start = self._key.seek_location
+        stop = start + self._key.num_bytes + self._key.compressed_bytes
+
+        self.write_anew(sink)
+
+        newloc = self._key.seek_location
+        file._move_tree(oldloc, newloc)
+
+        self._freesegments.release(start, stop)
+        sink.set_file_length(self._freesegments.fileheader.end)
+        sink.flush()
+
     def extend(self, file, sink, data):
+        if self._needs_relocation_before_extend:
+            # metadata_start/basket_metadata_start (computed structurally by
+            # _build_out(), see _load_existing_ttree) are only guaranteed
+            # correct for a tree actually laid out the way Uproot's own
+            # write_anew() lays one out. A tree loaded from an existing file
+            # may have been written by ROOT instead, with a byte-for-byte
+            # different (but semantically equivalent) layout Uproot's writer
+            # doesn't replicate -- write_updates() patching at Uproot's
+            # assumed offsets would then silently corrupt unrelated bytes.
+            # write_anew() here establishes Uproot's own canonical layout for
+            # real (the same thing add_branches() already always does, which
+            # is why it already works on ROOT-written files), so every
+            # write_updates() patch from here on is guaranteed to target the
+            # right bytes. Costs one full blob rewrite, but only once, on the
+            # first extend() after loading -- not on every call.
+            self._relocate(file, sink)
+            self._needs_relocation_before_extend = False
+
         # expand capacity if this would REACH (not EXCEED) the existing capacity
         # that's because completely a full fBasketEntry has nowhere to put the
         # number of entries in the last basket (it's a fencepost principle thing),
@@ -479,17 +530,7 @@ def extend(self, file, sink, data):
                 datum["fBasketSeek"][: len(fBasketSeek)] = fBasketSeek
                 datum["fBasketEntry"][len(fBasketEntry)] = self._num_entries
 
-            oldloc = start = self._key.location
-            stop = start + self._key.num_bytes + self._key.compressed_bytes
-
-            self.write_anew(sink)
-
-            newloc = self._key.seek_location
-            file._move_tree(oldloc, newloc)
-
-            self._freesegments.release(start, stop)
-            sink.set_file_length(self._freesegments.fileheader.end)
-            sink.flush()
+            self._relocate(file, sink)
 
         provided = None
 
@@ -822,7 +863,24 @@ def extend(self, file, sink, data):
 
         self.write_updates(sink)
 
-    def write_anew(self, sink):
+    def _build_out(self):
+        """
+        Serializes this TTree's metadata blob (TTree + TBranches + TLeaves,
+        everything but the TBaskets themselves) into a list of byte chunks,
+        exactly as ``write_anew`` writes it.
+
+        As a side effect, this sets ``self._metadata_start`` and, for every
+        branch, ``datum["metadata_start"]``, ``datum["basket_metadata_start"]``,
+        and ``datum["tleaf_reference_number"]`` — the byte offsets (relative to
+        this blob) that ``write_updates`` writes at directly, without going
+        through ``write_anew`` again. This is the single source of truth for
+        that layout: computing it any other way (e.g. searching a decompressed
+        blob for the current value of a field) risks disagreeing with it,
+        especially when that field's current value is 0 and matches unrelated
+        bytes. ``uproot.writing.writable.WritableDirectory._load_existing_ttree``
+        calls this once, on a freshly reconstructed cascade, purely to recover
+        these offsets — the returned ``out`` is discarded, nothing is written.
+        """
         key_num_bytes = uproot.reading._key_format_big.size + 6
         name_asbytes = self._name.encode(errors="surrogateescape")
         title_asbytes = self._title.encode(errors="surrogateescape")
@@ -1200,6 +1258,10 @@ def write_anew(self, sink):
 
         self._metadata_start = sum(len(x) for x in out[:metadata_out_index])
 
+        return out
+
+    def write_anew(self, sink):
+        out = self._build_out()
         raw_data = b"".join(out)
         self._key = self._directory.add_object(
             sink,