Skip to content

Folds and Reductions

Page Maps

graph LR
  family["Python Programming"]
  program["Python Functional Programming"]
  section["Streaming Resilience Failure Handling"]
  page["Folds and Reductions"]
  capstone["Capstone evidence"]

  family --> program --> section --> page
  page -.applies in.-> capstone
flowchart LR
  orient["Orient on the page map"] --> read["Read the main claim and examples"]
  read --> inspect["Inspect the related code, proof, or capstone surface"]
  inspect --> verify["Run or review the verification path"]
  verify --> apply["Apply the idea back to the module and capstone"]

This lesson makes folds feel practical before they feel abstract. Treat a fold as the answer to a repeated-work problem: one traversal, one accumulator story, one place to reason about what is combined and when.

Start With the Repeated Traversal Smell

Once you can traverse a tree safely, the next temptation is to walk it again and again for each statistic you want. That is the waste this lesson surfaces.

  • If count, max depth, and text length are computed in separate passes, the design is repeating the same traversal logic.
  • If the accumulator shape is unclear, the fold reads like machinery instead of a summary of what the code is collecting.
  • If a streaming scan changes order or work bounds, the fold implementation has stopped matching the traversal contract from the previous lesson.

Core question:
How do you replace any structural-recursive aggregation with an iterative fold (catamorphism) that is stack-safe, fully lazy when needed, and capable of fusing arbitrary numbers of independent aggregations into a single O(N) traversal?

This lesson introduces folds as the disciplined aggregation layer on top of safe traversal:

  • use one traversal to compute several related summaries
  • make the accumulator explicit so you can see exactly what state is carried
  • preserve preorder and bounded-work behavior when moving from a simple fold to a streaming scan

The motivating tree statistics are useful because they are easy to explain and immediately show why fusion matters.

The naïve recursive solution is beautiful and obvious:

def recursive_stats(tree: TreeDoc) -> tuple[int, int, int]:
    count = 1
    length = len(tree.node.text)
    max_d = 0
    for child in tree.children:
        c, l, md = recursive_stats(child)
        count += c
        length += l
        max_d = max(max_d, md + 1)
    return count, length, max_d

It works perfectly as a specification of the aggregate… until the tree is 2000 levels deep and the recursion strategy itself becomes the weak point.

The production solution must still tell the same aggregation story, but it has to do so in a way that is iterative, fused, and optionally streaming.

That is what a fold gives us in this module: one explicit reduction over the tree where the recursion strategy is safe and the aggregation logic remains reviewable.

Use this when you routinely aggregate statistics over tree/document/graph structures and refuse to ship code that can RecursionError on pathological but legal inputs.

Outcome:
1. You will replace any recursive aggregation with an iterative fold that is formally terminating and stack-safe.
2. You will fuse arbitrary numbers of independent aggregations into a single traversal using immutable tuple accumulators.
3. You will ship streaming reductions (scan_tree) that are truly lazy and short-circuitable.

This section formalises exactly what you should be able to defend: termination, stack-safety, preorder consistency, fusion of aggregates, and bounded work for scans.


Concrete Motivating Example

Same deep Markdown-derived tree from M04C01:

graph TD
  root["Root (title)<br/>50 chars, depth 0"]
  s1["Section 1<br/>30 chars, depth 1"]
  s11["Subsection 1.1<br/>20 chars, depth 2"]
  s12["Subsection 1.2<br/>25 chars, depth 2"]
  s2["Section 2<br/>35 chars, depth 1"]
  s21["Subsection 2.1<br/>40 chars, depth 2"]
  leaf["Deep leaf<br/>10 chars, depth 2002"]
  root --> s1
  root --> s2
  s1 --> s11
  s1 --> s12
  s2 --> s21 --> leaf

Desired aggregates (computed in one pass):

total_nodes       = 2004
total_text_length = 85_050
max_depth         = 2002

We want all three numbers, plus optionally a running total after each node (for progress bars, early termination, etc.).


1. Laws & Invariants (machine-checked where possible)

All laws assume finite, acyclic TreeDoc inputs (always non-empty; root node exists).

Law Formal Statement Enforcement
Termination & Stack-Safety Completes in O(N) steps with O(1) call-stack frames for any finite acyclic tree. Formal proof via explicit stack + Hypothesis on 5000-node chains + CI recursion-limit guard.
Equivalence fold_tree(t, seed, f) == recursive_fold(t, seed, f) for all t (identical result and preorder application). Hypothesis test_fold_vs_recursive_equivalence.
Fusion Fused tuple fold equals separate folds: (count, length, max_d) == (fold_count(t), fold_len(t), fold_max_d(t)). Hypothesis test_fusion_equivalence.
Bounded-Work (scan_tree) Consuming first k partial accumulators visits exactly k nodes. Instrumented property test_scan_bounded_work.
Order Law Combiner applied in strict preorder (identical sequence to flatten(t) from M04C01). Property test test_fold_preorder_matches_flatten.

These laws turn “fold” from a pattern into a verifiable contract.


2. Decision Table – Which Fold Do You Actually Use?

Need Streaming Partials? Multiple Values? Recommended Variant
Single aggregate (tree) No No fold_tree or fold_tree_no_path
Multiple aggregates (tree) No Yes fold_tree_buffered with tuple accumulator (fused)
Running totals (tree) Yes No/Yes scan_tree (optionally with tuple accumulator)
Linear (list/iterator) No/Yes No linear_reduce / linear_accumulate

Never use recursive aggregation in library code.
Never run multiple separate folds when a single fused tuple fold gives the same result in one pass.


3. Public API Surface (end-of-Module-04 refactor note)

Refactor note: tree folds/scans live in funcpipe_rag.tree (capstone/src/funcpipe_rag/tree/folds.py).
funcpipe_rag.api.core re-exports the same names as a stable façade for the course modules.

from funcpipe_rag.api.core import (
    fold_count_length_maxdepth,
    fold_tree,
    fold_tree_buffered,
    fold_tree_no_path,
    linear_accumulate,
    linear_reduce,
    scan_count_length_maxdepth,
    scan_tree,
)

4. Reference Implementations

4.1 Recursive Specification (Didactic only)

def recursive_fold(
    tree: TreeDoc,
    seed: R,
    combiner: Callable[[R, TreeDoc, int, Path], R],
    *,
    depth: int = 0,
    path: Path = (),
) -> R:
    acc = combiner(seed, tree, depth, path)
    for i, child in enumerate(tree.children):
        acc = recursive_fold(child, acc, combiner, depth=depth + 1, path=path + (i,))
    return acc

4.2 Simple Explicit-Stack Fold (Readable reference)

def fold_tree(
    tree: TreeDoc,
    seed: R,
    combiner: Callable[[R, TreeDoc, int, Path], R],
) -> R:
    acc = seed
    stack: deque[tuple[TreeDoc, int, Path, int]] = deque([(tree, 0, (), 0)])
    while stack:
        node, depth, path, child_idx = stack.pop()
        if child_idx == 0:
            acc = combiner(acc, node, depth, path)
        if child_idx < len(node.children):
            stack.append((node, depth, path, child_idx + 1))
            child = node.children[child_idx]
            stack.append((child, depth + 1, path + (child_idx,), 0))
    return acc

4.3 Production Winner – Buffered-Path Fold (zero extra tuple allocation)

def fold_tree_buffered(
    tree: TreeDoc,
    seed: R,
    combiner: Callable[[R, TreeDoc, int, Path], R],
) -> R:
    """
    Same semantics as fold_tree but maintains the path using a single mutable list
    (no tuples on the traversal stack). One tuple per node is still created when
    calling combiner (intrinsic to passing the path).
    """
    acc = seed
    stack: deque[tuple[TreeDoc, int, int | None]] = deque([(tree, 0, None)])
    path: list[int] = []
    last_depth = 0

    while stack:
        node, depth, sib_idx = stack.pop()

        # Maintain mutable path prefix (identical logic to iter_flatten_buffered)
        if depth < last_depth:
            del path[depth:]
        if sib_idx is not None:
            if depth > len(path):
                path.append(sib_idx)
            else:
                path[depth-1] = sib_idx
        last_depth = depth

        acc = combiner(acc, node, depth, tuple(path[:depth]))

        for i in range(len(node.children)-1, -1, -1):
            stack.append((node.children[i], depth + 1, i))

    return acc

4.4 Optimised Fold Without Path (when path unused)

def fold_tree_no_path(
    tree: TreeDoc,
    seed: R,
    combiner: Callable[[R, TreeDoc, int], R],
) -> R:
    """Third parameter is depth."""
    acc = seed
    stack: deque[tuple[TreeDoc, int, int]] = deque([(tree, 0, 0)])
    while stack:
        node, depth, child_idx = stack.pop()
        if child_idx == 0:
            acc = combiner(acc, node, depth)
        if child_idx < len(node.children):
            stack.append((node, depth, child_idx + 1))
            stack.append((node.children[child_idx], depth + 1, 0))
    return acc

4.5 Streaming Scan (Running Totals – Truly Lazy)

def scan_tree(
    tree: TreeDoc,
    seed: R,
    combiner: Callable[[R, TreeDoc, int, Path], R],
) -> Iterator[R]:
    """Yield running accumulator after each node in preorder – O(k) work for first k yields."""
    acc = seed
    stack: deque[tuple[TreeDoc, int, Path, int]] = deque([(tree, 0, (), 0)])
    while stack:
        node, depth, path, child_idx = stack.pop()
        if child_idx == 0:
            acc = combiner(acc, node, depth, path)
            yield acc
        if child_idx < len(node.children):
            stack.append((node, depth, path, child_idx + 1))
            child = node.children[child_idx]
            stack.append((child, depth + 1, path + (child_idx,), 0))

Note: linear_accumulate (via itertools.accumulate) yields the initial seed as the first value; scan_tree yields only post-node accumulators (no initial seed yield).

4.6 Fused Multi-Value Example (Count + Length + Max Depth)

def fold_count_length_maxdepth(tree: TreeDoc) -> Tuple[int, int, int]:
    def step(acc: Tuple[int, int, int], tree: TreeDoc, depth: int, path: Path) -> Tuple[int, int, int]:
        count, length, max_d = acc
        return (
            count + 1,
            length + len(tree.node.text),
            max(max_d, depth)
        )
    return fold_tree_buffered(tree, (0, 0, 0), step)

5. Property-Based Proofs (capstone/tests/test_tree_folds.py)

@given(tree=tree_strategy())
def test_fold_vs_recursive_equivalence(tree):
    rec = recursive_fold(tree, (0, 0, 0), step_count_len_maxd)
    buf = fold_tree_buffered(tree, (0, 0, 0), step_count_len_maxd)
    assert rec == buf

@given(tree=tree_strategy())
def test_fusion_equivalence(tree):
    fused = fold_count_length_maxdepth(tree)
    count = fold_tree_no_path(tree, 0, lambda a, n, d: a + 1)
    length = fold_tree_no_path(tree, 0, lambda a, n, d: a + len(n.node.text))
    max_d = fold_tree_no_path(tree, 0, lambda a, n, d: max(a, d))
    assert fused == (count, length, max_d)

@given(tree=tree_strategy())
def test_fold_preorder_matches_flatten(tree):
    from funcpipe_rag.api.core import flatten
    order_via_fold: list[Path] = []
    fold_tree(tree, None, lambda _, n, d, p: order_via_fold.append(p))
    order_via_flatten = [c.metadata["path"] for c in flatten(tree)]
    assert order_via_fold == order_via_flatten

@given(tree=tree_strategy())
def test_fold_buffered_order_matches_simple(tree):
    order_simple: list[Path] = []
    order_buf: list[Path] = []
    fold_tree(tree, None, lambda _, n, d, p: order_simple.append(p))
    fold_tree_buffered(tree, None, lambda _, n, d, p: order_buf.append(p))
    assert order_simple == order_buf

6. Big-O & Allocation Guarantees (peak auxiliary memory)

Variant Time Call-stack Peak auxiliary heap Total allocations
fold_tree / scan_tree O(N) O(1) O(depth) O(N×depth) paths
fold_tree_buffered O(N) O(1) O(depth) O(N) paths (only on combine)
fold_tree_no_path O(N) O(1) O(depth) Zero paths

Result metadata (paths when used) is intrinsic; auxiliary overhead is only the explicit stack + one mutable path list.


7. Anti-Patterns & Immediate Fixes

Anti-Pattern Symptom Fix
Recursive aggregation in library code RecursionError on deep trees Replace with fold_tree_buffered
Separate folds for related stats 3–10× slower on large trees Fuse with tuple accumulator
Mutable accumulator Aliasing / nondeterminism Use immutable tuples
String concatenation in combiner Quadratic time Count lengths, join once at end

Fold the evidence the RAG run actually needs

Traversal answers “which node comes next?” A fold answers “what observation do we carry forward?” FuncPipe needs a document summary before choosing resource and failure policies: node count estimates embedding work, text length estimates input volume, and maximum depth helps diagnose pathological parser output.

Three separate folds are easy to read:

tree --count pass------> node count
tree --length pass-----> character count
tree --depth pass------> maximum depth

A fused fold visits each node once and carries a product accumulator:

tree --(count, length, maximum depth) pass--> all three observations

For the learning tree, trace the accumulator by hand:

Node Incoming accumulator Outgoing accumulator
root (0, 0, 0) (1, 4, 0)
a (1, 4, 0) (2, 5, 1)
b (2, 5, 1) (3, 6, 1)
c (3, 6, 1) (4, 7, 2)

Then run:

course=programs/python-programming/python-functional-programming
state="$course/capstone/module-reference-states/module-04"
venv=artifacts/venv/python-programming/python-functional-programming/capstone

PYTHONPATH="$state/src" \
  "$venv/bin/pytest" -q \
  "$state/tests/learning/test_module_04_resilient_streaming.py" \
  -k fused_tree_fold

The equality in that test is the preservation claim: fusion changes traversal count, not the three answers. It does not measure a speedup. On a small tree, tuple construction and a more complex combiner may erase the benefit. Measure when performance is the reason for fusion.

Choose the fold surface by required evidence:

Need Route Cost to make visible
One final value and no path fold_tree_no_path No path tuples
One final value with section provenance fold_tree_buffered One path tuple per combiner call
Progress after every node scan_tree One emitted accumulator per node
A readable oracle for shallow generated trees Recursive test fold Call-stack growth

The iterative folds assume a valid finite tree. Unlike flatten, the Module 04 fold implementation does not detect repeated object identities. Validate the tree at the boundary when untrusted construction can create cycles, or the fold may not terminate. “Stack safe” does not mean “safe for every graph.”

8. Pre-Core Quiz

  1. What does fusion preserve? → The observations computed from the same traversal order.
  2. What does fusion change? → The number of traversals and the shape of the accumulator.
  3. When should you use scan_tree? → When downstream code needs each running state rather than only the final state.
  4. Which route avoids path construction? → fold_tree_no_path.
  5. Does an explicit stack detect cycles automatically? → No; that is a separate input-validity contract.

9. Post-Core Exercise

Extend the fused accumulator with the number of empty text nodes.

  • Keep the original count, length, and depth results unchanged.
  • Write the expected four-tuple for the learning tree before editing.
  • Compare the fused answer with four separate fold_tree_no_path calls.
  • Explain why adding a mutable set to the accumulator would make review harder, even if the final count happened to be correct.

Move forward when you can distinguish an equivalence proof from a performance measurement.

Continue with: Memoization

The next core applies the same “preserve the answer, change the work” reasoning to repeated deterministic computation.