// Copyright 2026 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/// A tree data structure that backs the `immut/vector`.
///
/// # Representation contract
///
/// The shape invariants live on `Tree` in `types.mbt`; this is why they are
/// what they are. `invariants_wbtest.mbt` checks them mechanically, and is the
/// place to extend if you add a shape.
///
/// ## Two ways down
///
/// Indexing has a fast path and a general one, and `sizes` selects between
/// them:
///
/// - `None` — *radix descent* (`get_radix` / `set_radix`). The child index at
/// each level is read straight out of the index's bits,
/// `(index >> shift) & BITMASK`. No arithmetic on the way down, no search.
/// - `Some(cumulative)` — *search descent*. `get_branch_index` looks the child
/// up in the cumulative sizes, and the index is rebased for the recursion.
///
/// `get_radix` **aborts** the moment it meets a node carrying sizes: it has
/// already consumed the bits for that level, so it cannot switch strategies.
/// That is the whole reason `None` has to mean "this subtree is radix all the
/// way down" and not merely "this node is radix" — the property is inherited
/// by every reader, so it must hold transitively.
///
/// ## Full and radix are independent
///
/// Neither implies the other, and conflating them has bitten this file twice:
///
/// - *Radix does not imply full.* A radix node may have a partial **last**
/// child; that is exactly what `from_leaves` builds for a vector whose size
/// is not a power of 32. Radix indexing still works, because only the
/// children to the left of the target need to be full.
/// - *Full does not imply radix.* `append_right_leaf` and `push_end` grow a
/// relaxed node by patching its sizes array (`append_sizes_last`,
/// `push_sizes_last`, `update_sizes_last`) instead of re-deriving it, so a
/// node can reach exactly its full size while still carrying sizes. Marking
/// its parent radix on the strength of the child's *size* is what made `at`
/// abort in #4083; `compute_sizes` now also requires `Tree::is_radix`.
///
/// ## Who may produce `None`
///
/// Keeping the transitive property true means every construction site has to
/// earn it. There are only six:
///
/// - `compute_sizes` — the general case; requires every child full **and**
/// radix.
/// - `new_branch_left` — only for a full leaf, and it builds the whole spine.
/// - `radix_append_sizes` — only when a full leaf was appended, and only from
/// an already-radix node.
/// - `from_leaves` — builds nothing but radix nodes, top to bottom.
/// - `append_right_leaf`'s fresh root — only when `self` is already radix and
/// the appended leaf is full; it checks both inline instead of going through
/// `radix_append_sizes`.
/// - `rebalance`'s non-`top` wrapper — the one unchecked site. It is safe only
/// because the wrapper never escapes: the calling frame feeds it straight to
/// `tri_merge`, which keeps the child and drops the wrapper. See the comment
/// there.
///
/// ## Height comes from the caller
///
/// Nothing in a `Tree` records its own height; every walk is handed a `shift`
/// and decrements it by `NUM_BITS` per level. So a subtree's height is fixed by
/// where it is attached, and a builder has to place leaves at the depth the
/// capacity implies rather than at the depth the remaining data suggests —
/// attaching a short remainder one level too high is #4084. `Tree::size` makes
/// the same assumption for radix nodes, so a misplaced leaf corrupts sizes too.
//-----------------------------------------------------------------------------
// Hyperparameters
//-----------------------------------------------------------------------------
///|
/// The controlling factor of the tree depth.
/// This is the only parameter you normally would adjust.
const NUM_BITS = 5
///|
/// Invariant: `BRANCHING_FACTOR` is a power of 2.
const BRANCHING_FACTOR : Int = 1 << NUM_BITS
///|
const BITMASK : Int = BRANCHING_FACTOR - 1
///|
/// The threshold for switching to a linear search.
const LINEAR_THRESHOLD : Int = 4
///|
/// The $e_{max}$ parameter of the search step invariant.
const E_MAX : Int = 2
///|
/// $e_{max} / 2$.
let e_max_2 : Int = E_MAX / 2
//-----------------------------------------------------------------------------
// Constructors
//-----------------------------------------------------------------------------
///|
fn[T] Tree::empty() -> Tree[T] {
Empty
}
///|
/// Create a new tree with a single leaf. Note that the resulting tree is a left-skewed tree.
#owned(leaf)
fn[T] new_branch_left(leaf : FixedArray[T], shift : Int) -> Tree[T] {
// A full leaf keeps the branch radix-indexable (`None`). A partial leaf must
// be relaxed (`Some`): on its own it sits at the right edge, but once another
// leaf is appended to its right it becomes an interior, non-full leaf and
// radix indexing would read out of bounds.
let sizes : FixedArray[Int]? = if leaf.length() == BRANCHING_FACTOR {
None
} else {
Some([leaf.length()])
}
match shift {
0 => Leaf(leaf)
s => Node([new_branch_left(leaf, s - NUM_BITS)], sizes)
}
}
//-----------------------------------------------------------------------------
// Getters
//-----------------------------------------------------------------------------
///|
/// Get the element at the given index.
///
/// Precondition:
/// - `self` is of height `shift / NUM_BITS`.
fn[T] Tree::get(self : Tree[T], index : Int, shift : Int) -> T {
fn get_radix(node : Tree[T], shift : Int) -> T {
for cur = node, s = shift {
match cur {
Leaf(leaf) => break leaf[index & BITMASK]
Node(node, None) =>
continue node[radix_indexing(index, s)], s - NUM_BITS
Node(_, Some(_)) =>
abort("Unreachable: Node should not have sizes in get_radix")
Empty => abort("Index out of bounds")
}
}
}
match self {
Leaf(leaf) => leaf[index]
Node(children, Some(sizes)) => {
let branch_index = get_branch_index(sizes, index)
let sub_index = if branch_index == 0 {
index
} else {
index - sizes[branch_index - 1]
}
children[branch_index].get(sub_index, shift - NUM_BITS)
}
Node(_, None) => get_radix(self, shift)
Empty => abort("Index out of bounds")
}
}
//-----------------------------------------------------------------------------
// Mutators
//-----------------------------------------------------------------------------
///|
/// Set a value at the given index.
///
/// Precondition:
/// - `self` is of height `shift / NUM_BITS`.
#owned(value)
fn[T] Tree::set(self : Tree[T], index : Int, shift : Int, value : T) -> Tree[T] {
// TODO: optimize this as loop
fn set_radix(node : Tree[T], shift : Int) -> Tree[T] {
match node {
Leaf(leaf) => Leaf(immutable_set(leaf, index & BITMASK, value))
Node(node, None) => {
let sub_idx = radix_indexing(index, shift)
Node(
immutable_set(
node,
sub_idx,
set_radix(node[radix_indexing(index, shift)], shift - NUM_BITS),
),
None,
)
}
Node(_, Some(_)) =>
abort("Unreachable: Node should not have sizes in set_radix")
Empty => abort("Index out of bounds")
}
}
match self {
Leaf(leaf) => Leaf(immutable_set(leaf, index & BITMASK, value))
Node(children, Some(sizes)) => {
let branch_index = get_branch_index(sizes, index)
let sub_index = if branch_index == 0 {
index
} else {
index - sizes[branch_index - 1]
}
Node(
immutable_set(
children,
branch_index,
children[branch_index].set(sub_index, shift - NUM_BITS, value),
),
Some(sizes),
)
}
Node(_children, None) => set_radix(self, shift)
Empty => abort("Index out of bounds")
}
}
///|
/// Push a value to the end of the tree.
///
/// Precondition:
/// - The height of `self` = `shift` / `NUM_BITS` (the height starts from 0).
#owned(value)
fn[T] Tree::push_end(self : Tree[T], shift : Int, value : T) -> (Tree[T], Int) {
fn update_sizes_last(sizes : FixedArray[Int]?) -> FixedArray[Int]? {
match sizes {
Some(sizes) => {
let new_sizes = sizes.copy()
new_sizes[new_sizes.length() - 1] += 1
Some(new_sizes)
}
None => None
}
}
fn push_sizes_last(sizes : FixedArray[Int]?) -> FixedArray[Int]? {
match sizes {
Some(sizes) => Some(immutable_push(sizes, 1 + sizes[sizes.length() - 1]))
None => None
}
}
fn worker(node : Tree[T], shift : Int) -> Tree[T]? {
match node {
Leaf(leaf) => {
if shift != 0 {
abort(
"Unreachable: Leaf should not have a non-zero shift, which means we have not reached the bottom of the tree",
)
}
if leaf.length() < BRANCHING_FACTOR {
Some(Leaf(immutable_push(leaf, value)))
} else {
None
}
}
Node(nodes, sizes) => {
let len = nodes.length()
match worker(nodes[len - 1], shift - NUM_BITS) {
// We have successfully pushed the value, now duplicate its ancestor nodes.
// Pushing a single value always lands in a partial leaf, so a radix
// (`None`) node must switch to explicit sizes.
Some(new_node) => {
let new_nodes = nodes.copy()
new_nodes[len - 1] = new_node
let new_sizes = match sizes {
Some(_) => update_sizes_last(sizes)
None => compute_sizes(new_nodes, shift - NUM_BITS)
}
Some(Node(new_nodes, new_sizes))
}
// We need to create a new node to push the value.
None =>
if len < BRANCHING_FACTOR {
let new_nodes = immutable_push(
nodes,
new_branch_left([value], shift - NUM_BITS),
)
let new_sizes = match sizes {
Some(_) => push_sizes_last(sizes)
None => compute_sizes(new_nodes, shift - NUM_BITS)
}
Some(Node(new_nodes, new_sizes))
} else {
None
}
}
}
Empty => Some(Leaf([value]))
}
}
match worker(self, shift) {
Some(new_tree) => (new_tree, shift)
None => {
// The new root holds `self` plus a fresh single-element branch. That
// branch is always a partial leaf, so the root needs explicit sizes even
// when `self` was radix.
let new_branch = new_branch_left([value], shift)
let len = self.size(shift)
(
Node([self, new_branch], Some(FixedArray::from_array([len, 1 + len]))),
shift + NUM_BITS,
)
}
}
}
//-----------------------------------------------------------------------------
// Iteration
//-----------------------------------------------------------------------------
///|
/// For each element in the tree, apply the function `f`.
fn[A] Tree::each(self : Tree[A], f : (A) -> Unit raise?) -> Unit raise? {
match self {
Empty => ()
Leaf(l) => l.each(f)
Node(ns, _) => ns.each(t => t.each(f))
}
}
///|
fn[A : Eq] Tree::contains(self : Tree[A], value : A) -> Bool {
match self {
Empty => false
Leaf(elems) =>
for elem in elems {
guard elem != value else { break true }
} nobreak {
false
}
Node(children, _) =>
for child in children {
guard !child.contains(value) else { break true }
} nobreak {
false
}
}
}
///|
/// For each element in the tree, apply the function `f` with the index of the element.
fn[A] Tree::eachi(
self : Tree[A],
f : (Int, A) -> Unit raise?,
shift : Int,
start : Int,
) -> Unit raise? {
match self {
Empty => ()
Leaf(l) =>
for i, x in l {
f(start + i, x)
}
Node(ns, None) => {
let child_shift = shift - NUM_BITS
for child in ns; start = start {
child.eachi(f, child_shift, start)
continue start + (1 << shift)
}
}
Node(ns, Some(sizes)) => {
let child_shift = shift - NUM_BITS
for i, child in ns; offset = 0 {
child.eachi(f, child_shift, start + offset)
continue sizes[i]
}
}
}
}
///|
/// Fold the tree.
fn[A, B] Tree::fold(
self : Tree[A],
acc : B,
f : (B, A) -> B raise?,
) -> B raise? {
match self {
Empty => acc
Leaf(l) => l.fold(f, init=acc)
Node(n, _) => n.fold((acc, t) => t.fold(acc, f), init=acc)
}
}
///|
/// Fold the tree in reverse order.
fn[A, B] Tree::rev_fold(
self : Tree[A],
acc : B,
f : (B, A) -> B raise?,
) -> B raise? {
match self {
Empty => acc
Leaf(l) => l.rev_fold(f, init=acc)
Node(n, _) => n.rev_fold((acc, t) => t.rev_fold(acc, f), init=acc)
}
}
///|
/// Map the tree.
fn[A, B] Tree::map(self : Tree[A], f : (A) -> B raise?) -> Tree[B] raise? {
match self {
Empty => Empty
Leaf(l) => Leaf(l.map(f))
Node(n, szs) =>
Node(FixedArray::makei(n.length(), i => n[i].map(f)), copy_sizes(szs))
}
}
//-----------------------------------------------------------------------------
// Concatenation
//-----------------------------------------------------------------------------
///|
/// Concatenate two trees.
/// Should be called as with `top = true`.
///
/// Preconditions:
/// - `left` and `right` are not `Empty`.
/// - `left` and `right` are of height `left_shift / NUM_BITS` and `right_shift / NUM_BITS`, respectively.
fn[A] Tree::concat(
left : Tree[A],
left_shift : Int,
right : Tree[A],
right_shift : Int,
top : Bool,
) -> (Tree[A], Int) {
if left_shift > right_shift {
let (c, c_shift) = Tree::concat(
left.right_child(),
left_shift - NUM_BITS,
right,
right_shift,
false,
)
guard! c_shift == left_shift
return rebalance(left, c, Empty, left_shift, top)
} else if right_shift > left_shift {
let (c, c_shift) = Tree::concat(
left,
left_shift,
right.left_child(),
right_shift - NUM_BITS,
false,
)
guard! c_shift == right_shift
return rebalance(Empty, c, right, right_shift, top)
} else if left_shift == 0 {
// Handle Leaf case
let left_elems = left.leaf_elements()
let right_elems = right.leaf_elements()
let left_len = left_elems.length()
let right_len = right_elems.length()
let len = left_len + right_len
if top && len <= BRANCHING_FACTOR {
return (
Leaf(
FixedArray::makei(len, (i : Int) => {
if i < left_len {
left_elems[i]
} else {
right_elems[i - left_len]
}
}),
),
0,
)
} else {
return (
Node(
FixedArray::from_array([left, right]),
Some(FixedArray::from_array([left_len, len])),
),
NUM_BITS,
)
}
} else {
// Handle Node case
let (c, c_shift) = Tree::concat(
left.right_child(),
left_shift - NUM_BITS,
right.left_child(),
right_shift - NUM_BITS,
false,
)
guard! c_shift == left_shift
guard! c_shift == right_shift
return rebalance(left, c, right, left_shift, top)
}
}
///|
/// Given three `Node`s of the same height (`shift` / `NUM_BITS`), rebalance
/// them into as few nodes as their children need.
///
/// How wide can the merge get? `tri_merge` drops `left`'s last child and
/// `right`'s first, so it yields
/// `(left_arity - 1) + center_arity + (right_arity - 1)`, and `redis_plan` only
/// ever shortens that. With `left_arity` and `right_arity` at most
/// `BRANCHING_FACTOR`:
///
/// - from `Tree::concat`, whose leaf case hands over a two-child center, at
/// most `31 + 2 + 31 = 64` — two nodes;
/// - from `Tree::concat_with_suffix`, which splices the left vector's tail in
/// as a third leaf, at most `31 + 3 + 31 = 65` — three.
///
/// So the result is one node, or `ceil(n / BRANCHING_FACTOR)` of them, which is
/// never more than three. Assuming two was #4086: the leftover 33rd child made
/// a node wider than the radix index can address, and `radix_indexing` masked
/// child 32 back to 0.
///
/// `top` is `true` if the resulting node has no upper node.
/// Returns the new node and its shift.
fn[A] rebalance(
left : Tree[A],
center : Tree[A],
right : Tree[A],
shift : Int,
top : Bool,
) -> (Tree[A], Int) {
// Suppose H = shift / NUM_BITS
let t = tri_merge(left, center, right) // t is a list of trees of (H-1) height
let (nc, nc_len) = redis_plan(t)
let new_t = redis(t, nc, nc_len, shift - NUM_BITS) // new_t is a list of trees of (H-1) height
guard! new_t.length() == nc_len
if nc_len <= BRANCHING_FACTOR {
// All nodes can be accommodated in a single node
let node = Node(new_t, compute_sizes(new_t, shift - NUM_BITS)) // node of H height
if !top {
// This wrapper is the one place that puts `None` over a child without
// checking it, so it would break radix descent if it ever reached a
// finished tree. It cannot: `!top` means the caller is another
// `Tree::concat`/`concat_with_suffix` frame, whose `rebalance` feeds it
// straight to `tri_merge` and keeps only the child. The layer exists
// solely to match the height the caller asserts on.
return (Node(FixedArray::from_array([node]), None), shift + NUM_BITS)
// return (H+1) height node, add another layer to align with the case at the end of the thisfunction
} else {
return (node, shift)
// return H height node, no upper node so no need to add another layer on top of it
}
} else {
// Pack the redistributed children into as many nodes as they need.
//
// Two is not always enough: `Tree::concat`'s leaf case hands us a center of
// at most 2 children, so `tri_merge` yields at most 31 + 2 + 31 = 64 - but
// `Tree::concat_with_suffix` splices the left vector's tail in as a third
// leaf, and 31 + 3 + 31 = 65 leaves a 33rd child over. A node that wide
// breaks radix indexing, which masks the child index with `BITMASK`.
let count = (new_t.length() + BRANCHING_FACTOR - 1) / BRANCHING_FACTOR
let new_children = FixedArray::makei(count, i => {
let lo = i * BRANCHING_FACTOR
let hi = min(lo + BRANCHING_FACTOR, new_t.length())
let chunk = FixedArray::makei(hi - lo, j => new_t[lo + j])
Node(chunk, compute_sizes(chunk, shift - NUM_BITS)) // height H
})
return (
Node(new_children, compute_sizes(new_children, shift)),
shift + NUM_BITS,
) // return (H+1) height node
}
}
///|
/// Given three trees of the same height (if not `Empty`), merge them into one.
/// `left` and `right` might be `Node` or `Empty`.
/// `center` is always a `Node`.
/// The resulting array might be longer than `BRANCHING_FACTOR`,
/// which will be handled by `rebalance` later.
///
/// Preconditions:
/// - `left` and `right` are `Empty` or `Node`.
/// - `center` is `Node`.
///
/// Postconditions:
/// - The resulting array is of length `left.length() + center.length() + right.length()`.
/// - The height of a `Tree` in the resulting array is one less than the height of the input `Tree`s.
fn[A] tri_merge(
left : Tree[A],
center : Tree[A],
right : Tree[A],
) -> FixedArray[Tree[A]] {
if left.is_leaf() || !center.is_node() || right.is_leaf() {
abort("Unreachable: input to merge is invalid")
}
fn get_children(self : Tree[A]) -> FixedArray[Tree[A]] {
match self {
Node(children, _) => children
Empty => []
Leaf(_) => abort("Unreachable")
}
}
let left_children = get_children(left)
let center_children = get_children(center)
let right_children = get_children(right)
let left_len = left_children.length()
let left_len = if left_len == 0 { 0 } else { left_len - 1 }
let center_len = center_children.length()
let right_len = right_children.length()
let right_len = if right_len == 0 { 0 } else { right_len - 1 }
FixedArray::makei(left_len + center_len + right_len, i => {
if i < left_len {
left_children[i]
} else if i < left_len + center_len {
center_children[i - left_len]
} else if right_len > 0 {
right_children[1 + i - left_len - center_len]
} else {
abort("Unreachable")
}
})
}
///|
/// Create a redistribution plan for the tree.
///
/// Returns the per-node child counts and how many of them are live: the first
/// `new_len` entries of the returned array describe the new nodes, and
/// anything past that is scratch left over from the merge.
///
/// The plan preserves the total number of children and only ever shortens the
/// list, merging short nodes rightwards until the length is within `e_max_2`
/// of the optimum. Two things about the loop are worth spelling out, because
/// both look out of bounds and are not:
///
/// - `node_counts[i + 1]` (the carry step) can name the slot just past the live
/// prefix. It never does. The carry only advances while
/// `remaining + node_counts[i + 1] > BRANCHING_FACTOR`, so reaching the last
/// live slot with a carry would mean every slot before it holds a full
/// `BRANCHING_FACTOR`, hence `opt_len >= new_len - 1` — which contradicts the
/// `opt_len + e_max_2 < new_len` that let us into the loop at all.
/// - the leading scan `while node_counts[i] > BRANCHING_FACTOR - e_max_2` is
/// bounded for the same reason: if every live node were full there would be
/// no slack and the outer loop would not have been entered.
///
/// Termination: each pass drops `new_len` by one and `new_len` is bounded
/// below by `opt_len + e_max_2`. `i` is only rewound by one per pass, after
/// having advanced by at least one, so it never goes negative.
fn[A] redis_plan(t : FixedArray[Tree[A]]) -> (FixedArray[Int], Int) {
let node_counts = FixedArray::makei(t.length(), i => t[i].local_size())
let total_nodes = node_counts.fold(init=0, (acc, x) => acc + x)
// round up to the nearest integer of S/BRANCHING_FACTOR
let opt_len = (total_nodes + BRANCHING_FACTOR - 1) / BRANCHING_FACTOR
let mut new_len = t.length()
let mut i = 0
while opt_len + e_max_2 < new_len {
// Skip over all nodes satisfying the invariant.
while node_counts[i] > BRANCHING_FACTOR - e_max_2 {
i += 1
}
// Found short node, so redistribute over the next nodes
let mut remaining_nodes = node_counts[i]
while remaining_nodes > 0 {
let min_size = min(remaining_nodes + node_counts[i + 1], BRANCHING_FACTOR)
node_counts[i] = min_size
remaining_nodes = remaining_nodes + node_counts[i + 1] - min_size
i += 1
}
for j in i..<(new_len - 1) {
node_counts[j] = node_counts[j + 1]
}
new_len -= 1
i -= 1
}
return (node_counts, new_len)
}
///|
/// This function redistributes the nodes in `old_t` according to the plan in `node_counts`.
///
/// Preconditions:
/// - forall i in 0..node_nums, old_t[i] != Empty.
/// - `old_t` contains a list of trees, each of (`shift` / `NUM_BITS`) height.
/// - `node_counts` contains the number of children of each node in `new_t` (the redistributed version of `old_t`).
/// - `node_nums` is the number of live entries at the front of `node_counts` (`redis_plan`'s `new_len`);
/// `node_counts` itself may be longer, with scratch left over from the merge past that prefix.
///
/// Postcondition:
/// - The resulting trees in `new_t` are of the same height as trees in `old_t`.
///
/// The `guard! j < old_len` in each branch is a leftover: `old_t[j]` is already
/// indexed on the line above it, so a plan that overran would fault there
/// first and the guard can never fire. It is harmless because the only plans
/// that reach here come from `redis_plan`, which conserves the child total and
/// emits positive counts, so the cursor is exhausted exactly when the source
/// is. Do not read it as an active bounds check.
fn[A] redis(
old_t : FixedArray[Tree[A]],
node_counts : FixedArray[Int],
node_nums : Int,
shift : Int,
) -> FixedArray[Tree[A]] {
let old_len = old_t.length()
let new_t = FixedArray::make(node_nums, Empty)
let mut old_offset = 0
let mut j = 0 // the index of in the old tree
if shift == 0 {
// Handle Leaf case
let mut old_leaf_elems = ([] : FixedArray[_])
let mut old_leaf_len = 0
for i in 0.. Bool {
match self {
Leaf(_) => true
Node(_, None) => true
Node(_, Some(_)) => false
Empty => false
}
}
///|
/// Given a list of trees as `children` with heights of (`shift` / `NUM_BITS`), compute the sizes array of the subtrees.
fn[A] compute_sizes(
children : FixedArray[Tree[A]],
shift : Int,
) -> FixedArray[Int]? {
let len = children.length()
let sizes = FixedArray::make(len, 0)
let mut sum = 0
let mut flag = true
let full_subtree_size = BRANCHING_FACTOR << shift
for i in 0..)]>,
#| Some(),
#|)
),
)
@debug.debug_inspect(
Node([Empty, Leaf([42])], Some([0, 1])),
content=(
#|Node(
#| )]>,
#| Some(),
#|)
),
)
}