cow-bench.py

# SPDX-License-Identifier: Apache-2.0
#
# This is a standalone benchmark to evaluate copy-on-write in Python.
# The goal is to visit all the element of a tree by avoiding the branches that share the same name.
import random, uuid, timeit


# A random tree generator
def gen_name():
    return random.choice("abcdefghijklmnopqrstuvwxyz")


def gen_node(depth, max_depth):
    if depth > max_depth:
        yield (gen_name(), [])
    else:
        for child in range(random.randint(1, 3)):
            yield (gen_name() + gen_name() + gen_name(), gen_node(depth + 1, max_depth))


def gen_tree(max_depth):
    random.seed(10)
    return gen_node(0, max_depth).__next__()


# Helper to show a tree
def print_tree(node, indent):
    (name, childs) = node
    print(" " * indent + name)
    for child in childs:
        print_tree(child, indent + 1)


# print_tree(gen_tree(4), 0)


# Visitor implementation by mutating a "visited" set
def visit_mutate(visited, node):
    count = 0
    (name, childs) = node
    if name not in visited:
        visited.add(name)
        count += 1
        for child in childs:
            count += visit_mutate(visited, child)
    return count


# Pure visitor implementation using a copy-on-write pattern
def visit_cow(visited, node):
    visited = visited.copy()
    count = 0
    (name, childs) = node
    if name not in visited:
        visited.add(name)
        count += 1
        for child in childs:
            (new_count, new_visited) = visit_cow(visited, child)
            count += new_count
            visited = visited.union(new_visited)
    return (count, visited)


def visit_cow_wrapper(visited, node):
    return visit_cow(visited, node)[0]


# Benchmark helpers
def bench(impl):
    assert impl(set(), gen_tree(30)) == 7920


# mutate -  0.0151767710049171
# cow    -  4.07419129399932
print("mutate - ", timeit.timeit(lambda: bench(visit_mutate), number=1))
print("cow    - ", timeit.timeit(lambda: bench(visit_cow_wrapper), number=1))