java-topology/defects/root-cern-0003/test/test_root_cern_0003.py

"""
root-cern-0003: TTree::InitializeBranchLists fSeqBranches std::find O(B^2) MOAD-0001

Stubs the branch-list partitioning algorithm in pure Python to measure
the O(B^2) vs O(B) improvement from replacing std::find with an unordered_set.
"""
import time
import unittest


class FakeBranch:
    """Minimal stub for TBranch: identity is object identity (pointer equiv)."""
    def __init__(self, name, count_branch=None):
        self.name = name
        self.count_branch = count_branch  # None = no count leaf

    def __repr__(self):
        return f"Branch({self.name})"


def initialize_branch_lists_defective(branches, check_leaf_count=True):
    """
    Defective: std::find on seq_branches list for every branch.
    Returns (seq_branches, sorted_branches, comparisons).
    """
    seq_branches = []
    sorted_branches = []
    comparisons = 0

    def in_seq(b):
        nonlocal comparisons
        for s in seq_branches:
            comparisons += 1
            if s is b:
                return True
        return False

    if check_leaf_count:
        for branch in branches:
            if branch.count_branch is not None:
                if not in_seq(branch.count_branch):
                    seq_branches.append(branch.count_branch)

    for branch in branches:
        if not in_seq(branch):
            sorted_branches.append(branch)

    return seq_branches, sorted_branches, comparisons


def initialize_branch_lists_patched(branches, check_leaf_count=True):
    """
    Patched: unordered_set<TBranch*> mirror for O(1) lookup.
    Returns (seq_branches, sorted_branches, comparisons).
    """
    seq_branches = []
    seq_set = set()   # mirrors seq_branches -- O(1) lookup
    sorted_branches = []
    comparisons = 0

    if check_leaf_count:
        for branch in branches:
            if branch.count_branch is not None:
                comparisons += 1  # one hash lookup
                if branch.count_branch not in seq_set:
                    seq_branches.append(branch.count_branch)
                    seq_set.add(branch.count_branch)

    for branch in branches:
        comparisons += 1  # one hash lookup
        if branch not in seq_set:
            sorted_branches.append(branch)

    return seq_branches, sorted_branches, comparisons


def make_tree(n_branches, n_count):
    """
    Create n_branches FakeBranches.  The first n_count have a count_branch
    pointing to distinct count branches (which are among the n_branches set).
    """
    branches = [FakeBranch(f"b{i}") for i in range(n_branches)]
    count_pool = [FakeBranch(f"count_{i}") for i in range(n_count)]
    for i in range(min(n_count, n_branches)):
        branches[i].count_branch = count_pool[i % len(count_pool)]
    return branches


class TestRootCern0003(unittest.TestCase):

    def _run(self, n_branches, n_count):
        branches = make_tree(n_branches, n_count)

        seq_d, sort_d, cmp_d = initialize_branch_lists_defective(branches)
        seq_p, sort_p, cmp_p = initialize_branch_lists_patched(branches)

        # Functional equivalence: same seq and sorted sets (by name)
        self.assertEqual(
            sorted(b.name for b in seq_d),
            sorted(b.name for b in seq_p),
            "seq_branches mismatch between defective and patched"
        )
        self.assertEqual(
            sorted(b.name for b in sort_d),
            sorted(b.name for b in sort_p),
            "sorted_branches mismatch between defective and patched"
        )

        ratio = cmp_d / max(cmp_p, 1)
        return cmp_d, cmp_p, ratio

    def test_small_tree(self):
        """10 branches, 5 count leaves -- verify correctness."""
        cmp_d, cmp_p, ratio = self._run(10, 5)
        self.assertGreater(ratio, 1.0,
            f"Expected patched to use fewer comparisons: def={cmp_d}, pat={cmp_p}")

    def test_medium_tree(self):
        """100 branches, 50 count leaves -- should show clear speedup."""
        cmp_d, cmp_p, ratio = self._run(100, 50)
        # Defective: roughly O(B * S) = 100 * 50 = 5000+ comparisons
        # Patched: O(B) = 100+100 = 200 comparisons
        self.assertGreater(ratio, 5.0,
            f"Expected >5x speedup at B=100, S=50: got {ratio:.1f}x "
            f"(def={cmp_d}, pat={cmp_p})")

    def test_large_tree_cms_scale(self):
        """1000 branches, 500 count leaves -- CMS NanoAOD scale."""
        cmp_d, cmp_p, ratio = self._run(1000, 500)
        # Defective: ~O(1000 * 500) = 500000 comparisons
        # Patched: ~O(1000) = 1000 comparisons
        # Ratio: ~500x
        self.assertGreater(ratio, 100.0,
            f"Expected >100x speedup at B=1000: got {ratio:.1f}x "
            f"(def={cmp_d}, pat={cmp_p})")

    def test_worst_case_all_count(self):
        """All branches are count leaves -- maximum quadratic growth."""
        N = 200
        cmp_d, cmp_p, ratio = self._run(N, N)
        self.assertGreater(ratio, 50.0,
            f"Expected >50x speedup at N={N} all-count: got {ratio:.1f}x "
            f"(def={cmp_d}, pat={cmp_p})")

    def test_no_count_branches(self):
        """No count leaves -- degenerate case. Both variants are O(B).
        Defective does 0 comparisons (if-branch never taken for count lookup)
        but still scans fSeqBranches in the second loop. With zero count
        branches fSeqBranches is empty so the second scan is also 0. Patched
        does B hash lookups. Both are O(B) -- verify correctness only."""
        cmp_d, cmp_p, ratio = self._run(500, 0)
        # Both branches do O(B) work -- assert sorted_branches matches only
        branches = make_tree(500, 0)
        seq_d, sort_d, _ = initialize_branch_lists_defective(branches)
        seq_p, sort_p, _ = initialize_branch_lists_patched(branches)
        self.assertEqual(len(seq_d), 0, "No count branches means empty seq_d")
        self.assertEqual(len(seq_p), 0, "No count branches means empty seq_p")
        self.assertEqual(len(sort_d), len(sort_p),
            "sorted_branches length must match")

    def test_timing(self):
        """Wall-clock timing at B=500, S=500."""
        branches = make_tree(500, 500)

        t0 = time.perf_counter()
        for _ in range(10):
            initialize_branch_lists_defective(branches)
        dt_def = time.perf_counter() - t0

        t0 = time.perf_counter()
        for _ in range(10):
            initialize_branch_lists_patched(branches)
        dt_pat = time.perf_counter() - t0

        self.assertGreater(dt_def, dt_pat * 0.5,
            f"Patched ({dt_pat:.4f}s) should be faster than defective ({dt_def:.4f}s)")


if __name__ == "__main__":
    unittest.main(verbosity=2)