From 09237d5fb32ecb03051f57c34fa503e727fe49ff Mon Sep 17 00:00:00 2001 From: Kai Zhen <106396557+kaizhx@users.noreply.github.com> Date: Tue, 12 May 2026 07:04:10 +0200 Subject: [PATCH 1/4] Add initial persistent AVL tree map implementation --- .../collect/AbstractImmutableSortedMap.java | 2 + .../sosy_lab/common/collect/OurSortedMap.java | 1 + .../PathCopyingPersistentAvlTreeMap.java | 1478 +++++++++++++++++ 3 files changed, 1481 insertions(+) create mode 100644 src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java diff --git a/src/org/sosy_lab/common/collect/AbstractImmutableSortedMap.java b/src/org/sosy_lab/common/collect/AbstractImmutableSortedMap.java index a76378906..266c6d31e 100644 --- a/src/org/sosy_lab/common/collect/AbstractImmutableSortedMap.java +++ b/src/org/sosy_lab/common/collect/AbstractImmutableSortedMap.java @@ -22,6 +22,8 @@ abstract sealed class AbstractImmutableSortedMap extends AbstractImmutableMap implements OurSortedMap permits OurSortedMap.EmptyImmutableOurSortedMap, + PathCopyingPersistentAvlTreeMap, + PathCopyingPersistentAvlTreeMap.PartialSortedMap, PathCopyingPersistentTreeMap, PathCopyingPersistentTreeMap.PartialSortedMap { diff --git a/src/org/sosy_lab/common/collect/OurSortedMap.java b/src/org/sosy_lab/common/collect/OurSortedMap.java index 2cf416b12..0e9866384 100644 --- a/src/org/sosy_lab/common/collect/OurSortedMap.java +++ b/src/org/sosy_lab/common/collect/OurSortedMap.java @@ -27,6 +27,7 @@ sealed interface OurSortedMap extends NavigableMap permits AbstractImmutableSortedMap, DescendingSortedMap, + PathCopyingPersistentAvlTreeMap.PartialSortedMap, PathCopyingPersistentTreeMap.PartialSortedMap { Iterator> entryIterator(); diff --git a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java new file mode 100644 index 000000000..050173bb0 --- /dev/null +++ b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java @@ -0,0 +1,1478 @@ +// This file is part of SoSy-Lab Common, +// a library of useful utilities: +// https://github.com/sosy-lab/java-common-lib +// +// SPDX-FileCopyrightText: 2007-2020 Dirk Beyer +// +// SPDX-License-Identifier: Apache-2.0 + +package org.sosy_lab.common.collect; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Verify.verify; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.Iterators; +import com.google.common.collect.UnmodifiableIterator; +import com.google.errorprone.annotations.Immutable; +import com.google.errorprone.annotations.Var; +import com.google.errorprone.annotations.concurrent.LazyInit; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.io.Serializable; +import java.util.AbstractMap.SimpleImmutableEntry; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.Iterator; +import java.util.Map; +import java.util.NavigableMap; +import java.util.NavigableSet; +import java.util.Objects; +import java.util.SortedMap; +import java.util.function.BinaryOperator; +import java.util.function.Function; +import java.util.stream.Collector; +import java.util.stream.Collectors; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * This is an implementation of {@link PersistentSortedMap} that is based on the Adelson-Velsky and + * Landis (AVL) tree and path copying for copy-efficient persistency. + * + * @param The type of keys. + * @param The type of values. + */ +@Immutable(containerOf = {"K", "V"}) +@SuppressFBWarnings(value = "SING_SINGLETON_IMPLEMENTS_SERIALIZABLE", justification = "false alarm") +public final class PathCopyingPersistentAvlTreeMap< + K extends Comparable, V extends @Nullable Object> + extends AbstractImmutableSortedMap implements PersistentSortedMap, Serializable { + + private static final long serialVersionUID = 1041711151457528188L; + + @SuppressWarnings("unused") + @SuppressFBWarnings( + value = "EQ_DOESNT_OVERRIDE_EQUALS", + justification = "Inherits equals() according to specification.") + @Immutable(containerOf = {"K", "V"}) + private static final class Node + extends SimpleImmutableEntry { + + private static final long serialVersionUID = -7393505826652634501L; + + private final @Nullable Node left; + private final @Nullable Node right; + private final byte height; + + // Leaf node + Node(K pKey, V pValue) { + this(pKey, pValue, null, null); + } + + // Any node + Node(K pKey, V pValue, @Nullable Node pLeft, @Nullable Node pRight) { + super(pKey, pValue); + left = pLeft; + right = pRight; + height = (byte) (1 + Math.max(getHeight(left), getHeight(right))); + } + + boolean isLeaf() { + return left == null && right == null; + } + + static byte getHeight(@Nullable Node n) { + return n == null ? -1 : n.height; + } + + static byte getBalanceFactor(Node n) { + return (byte) (getHeight(n.left) - getHeight(n.right)); + } + + // Methods used for path copying + + @SuppressWarnings("ReferenceEquality") // cannot use equals() for check whether tree is the same + Node withLeftChild(Node newLeft) { + if (newLeft == left) { + return this; + } else { + return new Node<>(getKey(), getValue(), newLeft, right); + } + } + + @SuppressWarnings("ReferenceEquality") // cannot use equals() for check whether tree is the same + Node withRightChild(Node newRight) { + if (newRight == right) { + return this; + } else { + return new Node<>(getKey(), getValue(), left, newRight); + } + } + + static int countNodes(@Nullable Node n) { + if (n == null) { + return 0; + } + return countNodes(n.left) + 1 + countNodes(n.right); + } + } + + // static creation methods + + private static final PathCopyingPersistentAvlTreeMap EMPTY_MAP = + new PathCopyingPersistentAvlTreeMap(null); + + @SuppressWarnings("unchecked") + public static , V extends @Nullable Object> + PersistentSortedMap of() { + return (PersistentSortedMap) EMPTY_MAP; + } + + public static , V extends @Nullable Object> + PersistentSortedMap copyOf(Map map) { + checkNotNull(map); + + if (map instanceof PathCopyingPersistentAvlTreeMap) { + return (PathCopyingPersistentAvlTreeMap) map; + } + + @Var PersistentSortedMap result = of(); + for (Map.Entry entry : map.entrySet()) { + result = result.putAndCopy(entry.getKey(), entry.getValue()); + } + return result; + } + + /** + * Return a {@link Collector} that accumulates elements into a {@link + * PathCopyingPersistentAvlTreeMap}. Keys and values are the result of the respective functions. If + * duplicate keys appear, the collector throws an {@link IllegalArgumentException}. + */ + public static , V extends @Nullable Object> + Collector> toPathCopyingPersistentAvlTreeMap( + Function keyFunction, + Function valueFunction) { + return toPathCopyingPersistentAvlTreeMap( + keyFunction, + valueFunction, + (k, v) -> { + throw new IllegalArgumentException("Duplicate key " + k); + }); + } + + /** + * Return a {@link Collector} that accumulates elements into a {@link + * PathCopyingPersistentAvlTreeMap}. Keys and values are the result of the respective functions. + * Duplicate keys are resolved using the given merge function. + */ + public static , V extends @Nullable Object> + Collector> toPathCopyingPersistentAvlTreeMap( + Function keyFunction, + Function valueFunction, + BinaryOperator mergeFunction) { + checkNotNull(keyFunction); + checkNotNull(valueFunction); + checkNotNull(mergeFunction); + return Collectors.collectingAndThen( + Collectors.toMap( + keyFunction, + valueFunction, + mergeFunction, + () -> CopyOnWriteSortedMap.copyOf(PathCopyingPersistentAvlTreeMap.of())), + CopyOnWriteSortedMap::getSnapshot); + } + + // state and constructor + + private final @Nullable Node root; + + @SuppressWarnings("Immutable") + @LazyInit + private transient @Nullable NavigableSet> entrySet; + + @LazyInit private transient int size; + + private PathCopyingPersistentAvlTreeMap(@Nullable Node pRoot) { + root = pRoot; + } + + // private utility methods + + @SuppressWarnings("unchecked") + private static , V> @Nullable Node findNode( + Object key, Node root) { + checkNotNull(key); + return findNode((K) key, root); + } + + private static , V> @Nullable Node findNode( + K key, Node root) { + checkNotNull(key); + + @Var Node current = root; + while (current != null) { + int comp = key.compareTo(current.getKey()); + + if (comp < 0) { + // key < current.data + + current = current.left; + + } else if (comp > 0) { + // key > current.data + + current = current.right; + + } else { + // key == current.data + + return current; + } + } + return null; + } + + /** + * Find the node with the smallest key in a given non-empty subtree. + * + * @param root The subtree to search in. + * @return The node with the smallest key. + * @throws NullPointerException If tree is empty. + */ + private static , V> Node findSmallestNode(Node root) { + @Var Node current = root; + while (current.left != null) { + current = current.left; + } + return current; + } + + /** + * Find the node with the largest key in a given non-empty subtree. + * + * @param root The subtree to search in. + * @return The node with the largest key. + * @throws NullPointerException If tree is empty. + */ + private static , V> Node findLargestNode(Node root) { + @Var Node current = root; + while (current.right != null) { + current = current.right; + } + return current; + } + + /** + * Given a key and a tree, find the node in the tree with the given key, or (if {@code inclusive} + * is false or there is no such node) the node with the smallest key that is still greater than + * the key to look for. In terms of {@link NavigableMap} operations, this returns the node for + * {@code map.tailMap(key, inclusive).first()}. Returns null if the tree is empty or there is no + * node that matches (i.e., key is larger than the largest key in the map). + * + * @param key The key to search for. + * @param root The tree to look in. + * @return A node or null. + */ + private static , V> @Nullable Node findNextGreaterNode( + K key, Node root, boolean inclusive) { + checkNotNull(key); + + @Var Node result = null; // this is always greater than key + + @Var Node current = root; + while (current != null) { + int comp = key.compareTo(current.getKey()); + + if (comp < 0) { + // key < current.data + // All nodes to the right of current are irrelevant because they are too big. + // current is the best candidate we have found so far, so it becomes the new result + // (current is always smaller than the previous result and still bigger than key). + + result = current; + current = current.left; + + } else if (comp > 0) { + // key > current.data + // All nodes to the left of current are irrelevant because they are too small. + // current itself is too small, too. + + current = current.right; + + } else { + // key == current.data + + if (inclusive) { + return current; + } else { + // All nodes to the left of current are irrelevant because they are too small. + // current itself is too small, too. + // The left-most node in the right subtree of child is the result. + if (current.right == null) { + // no node smaller than key in this subtree + return result; + } else { + return findSmallestNode(current.right); + } + } + } + + if (current == null) { + // We have reached a leaf without finding the element. + return result; + } + } + return null; + } + + /** + * Given a key and a tree, find the node in the tree with the given key, or (if {@code inclusive} + * is false or there is no such node) the node with the largest key that is still smaller than the + * key to look for. In terms of {@link NavigableMap} operations, this returns the node for {@code + * map.headMap(key, inclusive).last()}. Returns null if the tree is empty or there is no node that + * matches (i.e., key is smaller than or equal to the smallest key in the map). + * + * @param key The key to search for. + * @param root The tree to look in. + * @return A node or null. + */ + private static , V> @Nullable Node findNextSmallerNode( + K key, Node root, boolean inclusive) { + checkNotNull(key); + + @Var Node result = null; // this is always smaller than key + + @Var Node current = root; + while (current != null) { + int comp = key.compareTo(current.getKey()); + + if (comp < 0) { + // key < current.data + // All nodes to the right of current are irrelevant because they are too big. + // current itself is too big, too. + + current = current.left; + + } else if (comp > 0) { + // key > current.data + // All nodes to the left of current are irrelevant because they are too small. + // current is the best candidate we have found so far, so it becomes the new result + // (current is always bigger than the previous result and still smaller than key). + + result = current; + current = current.right; + + } else { + // key == current.data + + if (inclusive) { + return current; + } else { + // All nodes to the right of current are irrelevant because they are too big. + // current itself is too big, too. + // The right-most node in the left subtree of child is the result. + if (current.left == null) { + // no node smaller than key in this subtree + return result; + } else { + return findLargestNode(current.left); + } + } + } + + if (current == null) { + // We have reached a leaf without finding the element. + return result; + } + } + return null; + } + + private static > boolean exceedsLowerBound( + K pKey, K pLowerBound, boolean pLowerInclusive) { + if (pLowerInclusive) { + return pKey.compareTo(pLowerBound) < 0; + } else { + return pKey.compareTo(pLowerBound) <= 0; + } + } + + private static > boolean exceedsUpperBound( + K pKey, K pUpperBound, boolean pUpperInclusive) { + if (pUpperInclusive) { + return pKey.compareTo(pUpperBound) > 0; + } else { + return pKey.compareTo(pUpperBound) >= 0; + } + } + + private static , V> int checkAssertions( + @Nullable Node current) { + if (current == null) { + return -1; + } + + // check property of binary search tree + if (current.left != null) { + checkState( + current.getKey().compareTo(current.left.getKey()) > 0, + "Tree has left child that is not smaller."); + } + if (current.right != null) { + checkState( + current.getKey().compareTo(current.right.getKey()) < 0, + "Tree has right child that is not bigger."); + } + + // Check recursively + int leftHeight = checkAssertions(current.left); + int rightHeight = checkAssertions(current.right); + // Check AVL invariants + checkState(Math.abs(leftHeight - rightHeight) <= 1, + "Invalid balance factor."); + checkState(Node.getHeight(current) == 1 + Math.max(leftHeight, + rightHeight)); + + return Node.getHeight(current); + } + + /** + * Check the map for violation of its invariants. + * + * @throws IllegalStateException If any invariant is violated. + */ + @VisibleForTesting + @SuppressWarnings("CheckReturnValue") + void checkAssertions() { + checkAssertions(root); + } + + // modifying methods + + /** + * Create a map instance with a given root node. + * + * @param newRoot A node or null (meaning the empty tree). + * @return A map instance with the given tree. + */ + @SuppressWarnings("ReferenceEquality") // cannot use equals() for check whether tree is the same + private PersistentSortedMap mapFromTree(@Nullable Node newRoot) { + if (newRoot == root) { + return this; + } else if (newRoot == null) { + return of(); + } else { + return new PathCopyingPersistentAvlTreeMap<>(newRoot); + } + } + + @Override + public PersistentSortedMap putAndCopy(K key, V value) { + return mapFromTree(putAndCopy0(checkNotNull(key), value, root)); + } + + private static , V> Node putAndCopy0( + K key, V value, @Var @Nullable Node current) { + // + + if (current == null) { + return new Node<>(key, value); + } + + int comp = key.compareTo(current.getKey()); + if (comp < 0) { + // key < current.data + Node newLeft = putAndCopy0(key, value, current.left); + current = current.withLeftChild(newLeft); + + } else if (comp > 0) { + // key > current.data + Node newRight = putAndCopy0(key, value, current.right); + current = current.withRightChild(newRight); + + } else { + current = new Node<>(key, value, current.left, current.right); + } + + // restore invariants + return rebalance(current); + } + + @SuppressWarnings("unchecked") + @Override + public PersistentSortedMap removeAndCopy(Object key) { + if (isEmpty()) { + return this; + } + return mapFromTree(removeAndCopy0((K) checkNotNull(key), root)); + } + + private static , V> @Nullable Node removeAndCopy0( + K key, @Var Node current) { + // + + int comp = key.compareTo(current.getKey()); + + if (comp < 0) { + // Target key is less than current key + if (current.left == null) { + // Target key is not in map. + return current; + } + + // recursively descent left child of current node + Node newLeft = removeAndCopy0(key, current.left); + current = current.withLeftChild(newLeft); + + + } else { + // Target key is greater than or equal to current key + + if (comp > 0) { + // Target key is greater than current key + if (current.right == null) { + // Target key is not in map. + return current; + } + + // recursively descent right child of current node + Node newRight = removeAndCopy0(key, current.right); + current = current.withRightChild(newRight); + return rebalance(current); + } + + if (comp == 0) { + // Target key is equal to current key + + // Case: Current node has one or no children + + // This covers leaf nodes as well because then both current.left == null and + // current.right == null are true, returning null in the process + if (current.left == null) { + return current.right; + } + + if (current.right == null) { + return current.left; + } + + // Case: Current node has two children + + // Use Browns improved deletion algorithm: https://arxiv.org/abs/2406.05162 + // Choose the replacement node from the taller of the two subtrees instead of always + // choosing either successor or predecessor, which can reduce the amount of rebalancing. + if (Node.getBalanceFactor(current) > 0) { + // Left subtree is taller than right subtree + + // Find predecessor + Node predecessor = findLargestNode(current.left); + // Remove predecessor + Node newLeft = removeMaximumNodeInTree(current.left); + // Replace predecessor + current = new Node<>(predecessor.getKey(), predecessor.getValue(), newLeft, + current.right); + + } else { + // Right subtree is taller than left tree or both are same height + + // Find successor + Node successor = findSmallestNode(current.right); + // Remove successor + Node newRight = removeMinimumNodeInTree(current.right); + // Replace successor + current = new Node<>(successor.getKey(), successor.getValue(), current.left, + newRight); + } + } + } + return rebalance(current); + } + + /** + * Unconditionally delete the node with the smallest key in a given subtree. + * + * @return A new subtree reflecting the change. + */ + private static @Nullable Node removeMinimumNodeInTree(@Var Node current) { + if (current.left == null) { + // + return current.right; + } + + // recursive descent + Node newLeft = removeMinimumNodeInTree(current.left); + current = current.withLeftChild(newLeft); + + return rebalance(current); + } + + /** + * Unconditionally delete the node with the largest key in a given subtree. + * + * @return A new subtree reflecting the change. + */ + private static @Nullable Node removeMaximumNodeInTree(@Var Node current) { + if (current.right == null) { + // + return current.left; + } + + // recursive descent + Node newRight = removeMaximumNodeInTree(current.right); + current = current.withRightChild(newRight); + + return rebalance(current); + } + + + private static Node rebalance(Node current) { + if (Math.abs(Node.getBalanceFactor(current)) <= 1) { + return current; + } + + return restructure(current); + } + + private static Node restructure(Node current) { + // When an imbalance in an AVL tree is detected, rotations always rebalance a specific set of + // three connected nodes: (x,y,z) with z being the first unbalanced node, y the taller child + // of z and x the taller child of y. + // Trinode restructuring is a rebalancing operation that achieves the same structural + // result as the four standard AVL rotation cases by treating them as one unified action: + // Rename the nodes x,y,z as a,b and c based on their in-order traversal such that a < b < c. + // Make b the new parent of the local subtree, with a becoming the left child and c the + // right child of b, while attaching the four original subtrees T0,T1,T2,T3 to the leaves of + // a and c in their original relative order. + + Node z = current; + Node y; + Node x; + + Node a; + Node b; + Node c; + + // Subtrees of the three nodes involved + Node T0; + Node T1; + Node T2; + Node T3; + + // Check if node is left-heavy or right-heavy + if (Node.getBalanceFactor(z) > 1) { + // Node is left-heavy + if (Node.getBalanceFactor(z.left) >= 0) { + // LL case: + // x < y < z + y = z.left; + x = y.left; + + T0 = x.left; + T1 = x.right; + T2 = y.right; + T3 = z.right; + + a = x; + b = y; + c = z; + + } else { + // LR case: + // y < x < z + y = z.left; + x = y.right; + + T0 = y.left; + T1 = x.left; + T2 = x.right; + T3 = z.right; + + a = y; + b = x; + c = z; + } + } else { + // Node.getBalanceFactor(z) < -1 is true + // Node is right-heavy + if (Node.getBalanceFactor(z.right) >= 0) { + // RL case: + // z < x < y + y = z.right; + x = y.left; + + T0 = z.left; + T1 = x.left; + T2 = x.right; + T3 = y.right; + + a = z; + b = x; + c = y; + + } else { + // RR case: + // z < y < x + y = z.right; + x = y.right; + + T0 = z.left; + T1 = y.left; + T2 = x.left; + T3 = x.right; + + a = z; + b = y; + c = x; + } + } + + return rebuildTrinode(a, b, c, T0, T1, T2, T3); + } + + private static Node rebuildTrinode(Node a, Node b, Node c, + @Nullable Node T0, @Nullable Node T1, + @Nullable Node T2, @Nullable Node T3) { + Node newA = new Node<>(a.getKey(), a.getValue(), T0, T1); + Node newC = new Node<>(c.getKey(), c.getValue(), T2, T3); + + return new Node<>(b.getKey(), b.getValue(), newA, newC); + } + + // read operations + + @Override + @SuppressWarnings("ReferenceEquality") // comparing nodes with equals would not suffice + public boolean equals(@Nullable Object pObj) { + if (pObj instanceof PathCopyingPersistentAvlTreeMap + && ((PathCopyingPersistentAvlTreeMap) pObj).root == root) { + return true; + } + return super.equals(pObj); + } + + @Override + @SuppressWarnings("RedundantOverride") // to document that using super.hashCode is intended + public int hashCode() { + return super.hashCode(); + } + + @Override + public @Nullable Entry getEntry(Object pKey) { + return findNode(pKey, root); + } + + @Override + public Iterator> entryIterator() { + return EntryInOrderIterator.create(root); + } + + @Override + public Iterator> descendingEntryIterator() { + return DescendingEntryInOrderIterator.create(root); + } + + @Override + public PersistentSortedMap empty() { + return of(); + } + + @Override + public boolean containsKey(Object pObj) { + return findNode(pObj, root) != null; + } + + @Override + public @Nullable V get(Object pObj) { + Node node = findNode(pObj, root); + return node == null ? null : node.getValue(); + } + + @Override + public V getOrDefault(Object pKey, V pDefaultValue) { + Node node = findNode(pKey, root); + return node == null ? pDefaultValue : node.getValue(); + } + + @Override + public boolean isEmpty() { + return root == null; + } + + @Override + public int size() { + if (size <= 0) { + size = Node.countNodes(root); + } + return size; + } + + @Override + public @Nullable Entry firstEntry() { + if (isEmpty()) { + return null; + } + + return findSmallestNode(root); + } + + @Override + public @Nullable Entry lastEntry() { + if (isEmpty()) { + return null; + } + + return findLargestNode(root); + } + + @Override + public @Nullable Entry ceilingEntry(K pKey) { + return findNextGreaterNode(pKey, root, /* inclusive= */ true); + } + + @Override + public @Nullable Entry floorEntry(K pKey) { + return findNextSmallerNode(pKey, root, /* inclusive= */ true); + } + + @Override + public @Nullable Entry higherEntry(K pKey) { + return findNextGreaterNode(pKey, root, /* inclusive= */ false); + } + + @Override + public @Nullable Entry lowerEntry(K pKey) { + return findNextSmallerNode(pKey, root, /* inclusive= */ false); + } + + @Override + public @Nullable Comparator comparator() { + return null; + } + + @Override + public OurSortedMap descendingMap() { + return new DescendingSortedMap<>(this); + } + + @Override + public NavigableSet> entrySet() { + if (entrySet == null) { + entrySet = new SortedMapEntrySet<>(this); + } + return entrySet; + } + + @Override + public OurSortedMap subMap( + K pFromKey, boolean pFromInclusive, K pToKey, boolean pToInclusive) { + checkNotNull(pFromKey); + checkNotNull(pToKey); + + return PartialSortedMap.create(root, pFromKey, pFromInclusive, pToKey, pToInclusive); + } + + @Override + public OurSortedMap headMap(K pToKey, boolean pToInclusive) { + checkNotNull(pToKey); + + return PartialSortedMap.create( + root, null, /* pFromInclusive= */ true, pToKey, /* pToInclusive= */ pToInclusive); + } + + @Override + public OurSortedMap tailMap(K pFromKey, boolean pFromInclusive) { + checkNotNull(pFromKey); + + return PartialSortedMap.create( + root, pFromKey, /* pFromInclusive= */ pFromInclusive, null, /* pToInclusive= */ false); + } + + /** + * Tree iterator with in-order iteration returning node objects, with possibility for lower and + * upper bound. + * + * @param The type of keys. + * @param The type of values. + */ + private static final class EntryInOrderIterator< + K extends Comparable, V extends @Nullable Object> + extends UnmodifiableIterator> { + + // invariants: + // stack.top is always the next element to be returned + // (i.e., its left subtree has already been handled) + + private final Deque> stack; + + // If not null, iteration stops at this key. + private final @Nullable K highKey; + private final boolean highInclusive; // only relevant if highKey != null + + static , V> Iterator> create( + @Nullable Node root) { + if (root == null) { + return Collections.emptyIterator(); + } else { + return new EntryInOrderIterator<>( + root, null, /* pLowInclusive= */ false, null, /* pHighInclusive= */ false); + } + } + + /** + * Create a new iterator with an optional lower and upper bound. + * + * @param pFromKey null or inclusive lower bound + * @param pToKey null or exclusive lower bound + */ + static , V> Iterator> createWithBounds( + @Nullable Node root, + @Nullable K pFromKey, + boolean pFromInclusive, + @Nullable K pToKey, + boolean pToInclusive) { + if (root == null) { + return Collections.emptyIterator(); + } else { + return new EntryInOrderIterator<>(root, pFromKey, pFromInclusive, pToKey, pToInclusive); + } + } + + private EntryInOrderIterator( + Node root, + @Nullable K pLowKey, + boolean pLowInclusive, + @Nullable K pHighKey, + boolean pHighInclusive) { + stack = new ArrayDeque<>(); + highKey = pHighKey; + highInclusive = pHighInclusive; + + if (pLowKey == null) { + pushLeftMostNodesOnStack(root); + } else { + // TODO: optimize: this iterates twice through the tree + pushNodesToKeyOnStack(root, findNextGreaterNode(pLowKey, root, pLowInclusive).getKey()); + } + stopFurtherIterationIfOutOfRange(); + } + + @Override + public boolean hasNext() { + return !stack.isEmpty(); + } + + private void pushLeftMostNodesOnStack(@Var Node current) { + while (current.left != null) { + stack.push(current); + current = current.left; + } + stack.push(current); + } + + private void pushNodesToKeyOnStack(@Var Node current, K key) { + while (current != null) { + int comp = key.compareTo(current.getKey()); + + if (comp < 0) { + stack.push(current); + current = current.left; + + } else if (comp > 0) { + // This node and it's left subtree can be ignored completely. + current = current.right; + + } else { + stack.push(current); + return; + } + } + throw new AssertionError( + "PartialEntryInOrderIterator created with lower bound that is not in map"); + } + + private void stopFurtherIterationIfOutOfRange() { + if (highKey != null + && !stack.isEmpty() + && exceedsUpperBound(stack.peek().getKey(), highKey, highInclusive)) { + // We have reached the end, next element would already be too large + stack.clear(); + } + } + + @Override + public Map.Entry next() { + Node current = stack.pop(); + // this is the element to be returned + + // if it has a right subtree, + // push it on stack so that it will be handled next + if (current.right != null) { + pushLeftMostNodesOnStack(current.right); + } + + stopFurtherIterationIfOutOfRange(); + + return current; + } + } + + /** + * Reverse tree iterator with in-order iteration returning node objects, with possibility for + * lower and upper bound. The iteration starts at the upper bound and goes to the lower bound. + * + * @param The type of keys. + * @param The type of values. + */ + private static final class DescendingEntryInOrderIterator< + K extends Comparable, V extends @Nullable Object> + extends UnmodifiableIterator> { + + // invariants: + // stack.top is always the next element to be returned + // (i.e., its right subtree has already been handled) + + private final Deque> stack; + + // If not null, iteration stops at this key. + private final @Nullable K lowKey; + private final boolean lowInclusive; // only relevant if lowKey != null + + static , V> Iterator> create( + @Nullable Node root) { + if (root == null) { + return Collections.emptyIterator(); + } else { + return new DescendingEntryInOrderIterator<>( + root, null, /* pLowInclusive= */ false, null, /* pHighInclusive= */ false); + } + } + + /** + * Create a new iterator with an optional lower and upper bound. + * + * @param pFromKey null or inclusive lower bound + * @param pToKey null or exclusive lower bound + */ + static , V> Iterator> createWithBounds( + @Nullable Node root, + @Nullable K pFromKey, + boolean pFromInclusive, + @Nullable K pToKey, + boolean pToInclusive) { + if (root == null) { + return Collections.emptyIterator(); + } else { + return new DescendingEntryInOrderIterator<>( + root, pFromKey, pFromInclusive, pToKey, pToInclusive); + } + } + + private DescendingEntryInOrderIterator( + Node root, + @Nullable K pLowKey, + boolean pLowInclusive, + @Nullable K pHighKey, + boolean pHighInclusive) { + stack = new ArrayDeque<>(); + lowKey = pLowKey; + lowInclusive = pLowInclusive; + + if (pHighKey == null) { + pushRightMostNodesOnStack(root); + } else { + // TODO: optimize: this iterates twice through the tree + pushNodesToKeyOnStack(root, findNextSmallerNode(pHighKey, root, pHighInclusive).getKey()); + } + stopFurtherIterationIfOutOfRange(); + } + + @Override + public boolean hasNext() { + return !stack.isEmpty(); + } + + private void pushRightMostNodesOnStack(@Var Node current) { + while (current.right != null) { + stack.push(current); + current = current.right; + } + stack.push(current); + } + + private void pushNodesToKeyOnStack(@Var Node current, K key) { + while (current != null) { + int comp = key.compareTo(current.getKey()); + + if (comp > 0) { + stack.push(current); + current = current.right; + + } else if (comp < 0) { + // This node and it's right subtree can be ignored completely. + current = current.left; + + } else { + stack.push(current); + return; + } + } + throw new AssertionError( + "PartialEntryInOrderIterator created with upper bound that is not in map"); + } + + private void stopFurtherIterationIfOutOfRange() { + if (lowKey != null + && !stack.isEmpty() + && exceedsLowerBound(stack.peek().getKey(), lowKey, lowInclusive)) { + // We have reached the end, next element would already be too small + stack.clear(); + } + } + + @Override + public Map.Entry next() { + Node current = stack.pop(); + // this is the element to be returned + + // if it has a left subtree, + // push it on stack so that it will be handled next + if (current.left != null) { + pushRightMostNodesOnStack(current.left); + } + + stopFurtherIterationIfOutOfRange(); + + return current; + } + } + + /** + * Partial map implementation for {@link SortedMap#subMap(Object, Object)} etc. At least one bound + * (upper/lower) needs to be present. The range needs to contain at least one mapping. + * + * @param The type of keys. + * @param The type of values. + */ + @Immutable(containerOf = {"K", "V"}) + static final class PartialSortedMap, V extends @Nullable Object> + extends AbstractImmutableSortedMap implements OurSortedMap, Serializable { + + static , V> OurSortedMap create( + Node pRoot, + @Nullable K pFromKey, + boolean pFromInclusive, + @Nullable K pToKey, + boolean pToInclusive) { + checkArgument(pFromKey != null || pToKey != null); + + if (pFromKey != null && pToKey != null) { + int comp = pFromKey.compareTo(pToKey); + if (comp == 0 && (!pFromInclusive || !pToInclusive)) { + return EmptyImmutableOurSortedMap.of(); + } + checkArgument(comp <= 0, "fromKey > toKey"); + } + + Node root = findBestRoot(pRoot, pFromKey, pFromInclusive, pToKey, pToInclusive); + if (root == null) { + return EmptyImmutableOurSortedMap.of(); + } + + @Var Node lowestNode = null; + if (pFromKey != null) { + lowestNode = findNextGreaterNode(pFromKey, root, pFromInclusive); + if (lowestNode == null) { + return EmptyImmutableOurSortedMap.of(); + } + } + + @Var Node highestNode = null; + if (pToKey != null) { + highestNode = findNextSmallerNode(pToKey, root, pToInclusive); + if (highestNode == null) { + return EmptyImmutableOurSortedMap.of(); + } + } + + if (pFromKey != null && pToKey != null) { + assert lowestNode != null && highestNode != null; + if (exceedsUpperBound(lowestNode.getKey(), pToKey, pToInclusive)) { + verify(exceedsLowerBound(highestNode.getKey(), pFromKey, pFromInclusive)); + + // no mappings in range + return EmptyImmutableOurSortedMap.of(); + } + } + + return new PartialSortedMap<>(root, pFromKey, pFromInclusive, pToKey, pToInclusive); + } + + // Find the best root for a given set of bounds + // (the lowest node in the tree that represents the complete range). + // Not using root directly but potentially only a subtree is more efficient. + private static , V> @Nullable Node findBestRoot( + @Nullable Node pRoot, + @Nullable K pFromKey, + boolean pFromInclusive, + @Nullable K pToKey, + boolean pToInclusive) { + + @Var Node current = pRoot; + while (current != null) { + + if (pFromKey != null && exceedsLowerBound(current.getKey(), pFromKey, pFromInclusive)) { + // current and left subtree can be ignored + current = current.right; + } else if (pToKey != null && exceedsUpperBound(current.getKey(), pToKey, pToInclusive)) { + // current and right subtree can be ignored + current = current.left; + } else { + // current is in range + return current; + } + } + + return null; // no mapping in range + } + + private static final long serialVersionUID = 5354023186935889223L; + + // Invariant: This map is never empty. + + private final Node root; + + // null if there is no according bound, in this case the "inclusive" boolean is irrelevant + @SuppressWarnings("serial") // This class only needs to be serializable if keys are. + private final @Nullable K fromKey; + + private final boolean fromInclusive; + + @SuppressWarnings("serial") // This class only needs to be serializable if keys are. + private final @Nullable K toKey; + + private final boolean toInclusive; + + @LazyInit private transient int size; + + @SuppressWarnings("Immutable") + @LazyInit + private transient @Nullable NavigableSet> entrySet; + + private PartialSortedMap( + Node pRoot, + @Nullable K pFromKey, + boolean pFromInclusive, + @Nullable K pToKey, + boolean pToInclusive) { + root = checkNotNull(pRoot); + fromKey = pFromKey; + fromInclusive = pFromInclusive; + toKey = pToKey; + toInclusive = pToInclusive; + + // check non-emptiness invariant + assert pFromKey == null + || containsKey(findNextGreaterNode(pFromKey, pRoot, fromInclusive).getKey()); + assert pToKey == null + || containsKey(findNextSmallerNode(pToKey, pRoot, toInclusive).getKey()); + } + + private boolean inRange(K key, boolean treatBoundsAsInclusive) { + return !tooLow(key, treatBoundsAsInclusive) && !tooHigh(key, treatBoundsAsInclusive); + } + + private boolean tooLow(K key, boolean treatBoundAsInclusive) { + return fromKey != null + && exceedsLowerBound(key, fromKey, treatBoundAsInclusive || fromInclusive); + } + + private boolean tooHigh(K key, boolean treatBoundAsInclusive) { + return toKey != null && exceedsUpperBound(key, toKey, treatBoundAsInclusive || toInclusive); + } + + @Override + public Iterator> entryIterator() { + return EntryInOrderIterator.createWithBounds( + root, fromKey, fromInclusive, toKey, toInclusive); + } + + @Override + public Iterator> descendingEntryIterator() { + return DescendingEntryInOrderIterator.createWithBounds( + root, fromKey, fromInclusive, toKey, toInclusive); + } + + @Override + @SuppressWarnings("ReferenceEquality") // comparing nodes with equals would not suffice + public boolean equals(@Nullable Object pObj) { + if (pObj instanceof PartialSortedMap other) { + if (root == other.root + && Objects.equals(fromKey, other.fromKey) + && fromInclusive == other.fromInclusive + && Objects.equals(toKey, other.toKey) + && toInclusive == other.toInclusive) { + return true; + } + } + return super.equals(pObj); + } + + @Override + @SuppressWarnings("RedundantOverride") // to document that using super.hashCode is intended + public int hashCode() { + return super.hashCode(); + } + + @Override + public boolean containsKey(Object pKey) { + @SuppressWarnings("unchecked") + K key = (K) checkNotNull(pKey); + return inRange(key, /* treatBoundsAsInclusive= */ false) && findNode(key, root) != null; + } + + @Override + public @Nullable Node getEntry(Object pKey) { + @SuppressWarnings("unchecked") + K key = (K) checkNotNull(pKey); + if (!inRange(key, /* treatBoundsAsInclusive= */ false)) { + return null; + } + Node node = findNode(key, root); + return node; + } + + @Override + public @Nullable V get(Object pKey) { + Node node = getEntry(pKey); + return node == null ? null : node.getValue(); + } + + @Override + public boolean isEmpty() { + return false; + } + + @Override + public int size() { + if (size == 0) { + size = Iterators.size(entryIterator()); + } + return size; + } + + @Override + public @Nullable Entry firstEntry() { + if (fromKey == null) { + return findSmallestNode(root); + } else { + return findNextGreaterNode(fromKey, root, fromInclusive); + } + } + + @Override + public @Nullable Entry lastEntry() { + if (toKey == null) { + return findLargestNode(root); + } else { + return findNextSmallerNode(toKey, root, toInclusive); + } + } + + @Override + public @Nullable Entry ceilingEntry(K pKey) { + Entry result = findNextGreaterNode(pKey, root, /* inclusive= */ true); + if (result != null && !inRange(result.getKey(), /* treatBoundsAsInclusive= */ false)) { + return null; + } + return result; + } + + @Override + public @Nullable Entry floorEntry(K pKey) { + Entry result = findNextSmallerNode(pKey, root, /* inclusive= */ true); + if (result != null && !inRange(result.getKey(), /* treatBoundsAsInclusive= */ false)) { + return null; + } + return result; + } + + @Override + public @Nullable Entry higherEntry(K pKey) { + Entry result = findNextGreaterNode(pKey, root, /* inclusive= */ false); + if (result != null && !inRange(result.getKey(), /* treatBoundsAsInclusive= */ false)) { + return null; + } + return result; + } + + @Override + public @Nullable Entry lowerEntry(K pKey) { + Entry result = findNextSmallerNode(pKey, root, /* inclusive= */ false); + if (result != null && !inRange(result.getKey(), /* treatBoundsAsInclusive= */ false)) { + return null; + } + return result; + } + + @Override + public @Nullable Comparator comparator() { + return null; + } + + @Override + public OurSortedMap descendingMap() { + return new DescendingSortedMap<>(this); + } + + @Override + public NavigableSet> entrySet() { + if (entrySet == null) { + entrySet = new SortedMapEntrySet<>(this); + } + return entrySet; + } + + @Override + public OurSortedMap subMap( + K pFromKey, boolean pFromInclusive, K pToKey, boolean pToInclusive) { + checkNotNull(pFromKey); + checkNotNull(pToKey); + + // If fromKey==pFromKey, we must forbid the combination of fromInclusive==false and + // pFromInclusive==true, because this would mean that the new range exceeds the old range. + // All other combinations of fromInclusive and pFromInclusive are allowed. + // So we accept equal keys if (!pFromInclusive || fromInclusive) holds, + // the latter being handled by inRange(). + checkArgument(inRange(pFromKey, !pFromInclusive)); + // Similarly for the upper bound. + checkArgument(inRange(pToKey, !pToInclusive)); + + return PartialSortedMap.create(root, pFromKey, pFromInclusive, pToKey, pToInclusive); + } + + @Override + public OurSortedMap headMap(K pToKey, boolean pInclusive) { + checkNotNull(pToKey); + checkArgument(inRange(pToKey, /* treatBoundsAsInclusive= */ !pInclusive)); + + return PartialSortedMap.create( + root, + fromKey, + /* pFromInclusive= */ fromInclusive, + pToKey, + /* pToInclusive= */ pInclusive); + } + + @Override + public OurSortedMap tailMap(K pFromKey, boolean pInclusive) { + checkNotNull(pFromKey); + checkArgument(inRange(pFromKey, /* treatBoundsAsInclusive= */ !pInclusive)); + + return PartialSortedMap.create( + root, pFromKey, /* pFromInclusive= */ pInclusive, toKey, /* pToInclusive= */ toInclusive); + } + } +} From c91ef07e29b8fb22c660b5a82efa4fecff0fac53 Mon Sep 17 00:00:00 2001 From: Kai Zhen <106396557+kaizhx@users.noreply.github.com> Date: Tue, 12 May 2026 07:18:33 +0200 Subject: [PATCH 2/4] Fix checkstyle --- .../PathCopyingPersistentAvlTreeMap.java | 50 +++++++++---------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java index 050173bb0..7c5b3ddb8 100644 --- a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java +++ b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java @@ -148,8 +148,8 @@ PersistentSortedMap copyOf(Map map) { /** * Return a {@link Collector} that accumulates elements into a {@link - * PathCopyingPersistentAvlTreeMap}. Keys and values are the result of the respective functions. If - * duplicate keys appear, the collector throws an {@link IllegalArgumentException}. + * PathCopyingPersistentAvlTreeMap}. Keys and values are the result of the respective functions. + * If duplicate keys appear, the collector throws an {@link IllegalArgumentException}. */ public static , V extends @Nullable Object> Collector> toPathCopyingPersistentAvlTreeMap( @@ -656,10 +656,10 @@ private static Node restructure(Node current) { Node c; // Subtrees of the three nodes involved - Node T0; - Node T1; - Node T2; - Node T3; + Node t0; + Node t1; + Node t2; + Node t3; // Check if node is left-heavy or right-heavy if (Node.getBalanceFactor(z) > 1) { @@ -670,10 +670,10 @@ private static Node restructure(Node current) { y = z.left; x = y.left; - T0 = x.left; - T1 = x.right; - T2 = y.right; - T3 = z.right; + t0 = x.left; + t1 = x.right; + t2 = y.right; + t3 = z.right; a = x; b = y; @@ -685,10 +685,10 @@ private static Node restructure(Node current) { y = z.left; x = y.right; - T0 = y.left; - T1 = x.left; - T2 = x.right; - T3 = z.right; + t0 = y.left; + t1 = x.left; + t2 = x.right; + t3 = z.right; a = y; b = x; @@ -703,10 +703,10 @@ private static Node restructure(Node current) { y = z.right; x = y.left; - T0 = z.left; - T1 = x.left; - T2 = x.right; - T3 = y.right; + t0 = z.left; + t1 = x.left; + t2 = x.right; + t3 = y.right; a = z; b = x; @@ -718,10 +718,10 @@ private static Node restructure(Node current) { y = z.right; x = y.right; - T0 = z.left; - T1 = y.left; - T2 = x.left; - T3 = x.right; + t0 = z.left; + t1 = y.left; + t2 = x.left; + t3 = x.right; a = z; b = y; @@ -729,14 +729,14 @@ private static Node restructure(Node current) { } } - return rebuildTrinode(a, b, c, T0, T1, T2, T3); + return rebuildTrinode(a, b, c, t0, t1, t2, t3); } private static Node rebuildTrinode(Node a, Node b, Node c, @Nullable Node T0, @Nullable Node T1, @Nullable Node T2, @Nullable Node T3) { - Node newA = new Node<>(a.getKey(), a.getValue(), T0, T1); - Node newC = new Node<>(c.getKey(), c.getValue(), T2, T3); + Node newA = new Node<>(a.getKey(), a.getValue(), T0, T1); + Node newC = new Node<>(c.getKey(), c.getValue(), T2, T3); return new Node<>(b.getKey(), b.getValue(), newA, newC); } From 279b052dd76470f14f2f5263dc9ce1b0c8c65380 Mon Sep 17 00:00:00 2001 From: Kai Zhen <106396557+kaizhx@users.noreply.github.com> Date: Tue, 12 May 2026 07:26:35 +0200 Subject: [PATCH 3/4] Fix checkstyle --- .../common/collect/PathCopyingPersistentAvlTreeMap.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java index 7c5b3ddb8..a2b105c00 100644 --- a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java +++ b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java @@ -733,10 +733,10 @@ private static Node restructure(Node current) { } private static Node rebuildTrinode(Node a, Node b, Node c, - @Nullable Node T0, @Nullable Node T1, - @Nullable Node T2, @Nullable Node T3) { - Node newA = new Node<>(a.getKey(), a.getValue(), T0, T1); - Node newC = new Node<>(c.getKey(), c.getValue(), T2, T3); + @Nullable Node t0, @Nullable Node t1, + @Nullable Node t2, @Nullable Node t3) { + Node newA = new Node<>(a.getKey(), a.getValue(), t0, t1); + Node newC = new Node<>(c.getKey(), c.getValue(), t2, t3); return new Node<>(b.getKey(), b.getValue(), newA, newC); } From 77e8fa9af90a7e8b1399268aa9e0b67b88df1190 Mon Sep 17 00:00:00 2001 From: Kai Zhen <106396557+kaizhx@users.noreply.github.com> Date: Tue, 12 May 2026 07:34:59 +0200 Subject: [PATCH 4/4] Fix checkstyle --- .../common/collect/PathCopyingPersistentAvlTreeMap.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java index a2b105c00..0ca6551f9 100644 --- a/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java +++ b/src/org/sosy_lab/common/collect/PathCopyingPersistentAvlTreeMap.java @@ -734,7 +734,8 @@ private static Node restructure(Node current) { private static Node rebuildTrinode(Node a, Node b, Node c, @Nullable Node t0, @Nullable Node t1, - @Nullable Node t2, @Nullable Node t3) { + @Nullable Node t2, + @Nullable Node t3) { Node newA = new Node<>(a.getKey(), a.getValue(), t0, t1); Node newC = new Node<>(c.getKey(), c.getValue(), t2, t3);