benchmark_accelerator.py

import sys
import os
import time
import random

# Add the project root to sys.path
sys.path.append(os.getcwd())

try:
    import treemap_accelerator as accelerator
    from spacehoarder import models
except ImportError:
    sys.path.append(os.path.dirname(os.getcwd()))
    import treemap_accelerator as accelerator
    from spacehoarder import models

# Always use mock classes for the benchmark to ensure consistency
class MockFileModel:
    def __init__(self, path, size, depth):
        self.path = path
        self.name = os.path.basename(path)
        self.size = size
        self.depth = depth
        self.contains = []

class MockDirModel(MockFileModel):
    def __init__(self, path, size, depth):
        super().__init__(path, size, depth)
        self.contains = []

def create_large_tree(depth=3, width=10):
    """Creates a tree with width^depth leaves roughly"""
    root = MockDirModel("/", 0, 0)
    
    def build_level(parent, current_depth):
        if current_depth >= depth:
            return
        
        for i in range(width):
            if current_depth == depth - 1:
                # Leaf
                f = MockFileModel(f"{parent.path}/f{i}", random.randint(100, 100000), current_depth + 1)
                parent.contains.append(f)
                parent.size += f.size
            else:
                # Dir
                d = MockDirModel(f"{parent.path}/d{i}", 0, current_depth + 1)
                build_level(d, current_depth + 1)
                parent.contains.append(d)
                parent.size += d.size

    build_level(root, 0)
    return root

print("Generating large tree...")
# 10^5 = 100,000 files if depth=5, width=10. Let's do depth=4, width=18 (~100k)
# Actually, depth=5, width=10 is 100,000 files at leaf level.
root_model = create_large_tree(depth=5, width=10)

def count_items(node):
    c = 1
    if hasattr(node, 'contains'):
        for child in node.contains:
            c += count_items(child)
    return c

count = count_items(root_model)
print(f"Generated tree with {count} items.")

print("Running benchmark WITH CACHING...")

# 1. Caching Step (Happens once on load)
t0 = time.time()
rust_cache = accelerator.RustLayoutModel(root_model)
t1 = time.time()
print(f"Cache creation (load time): {(t1-t0)*1000:.4f} ms")

# 2. Layout Step (Happens every resize)
num_colors = 10
start = time.time()
# width, height, min_size, pad, num_colors
# calculate_layout now expects the wrapper
rects = accelerator.calculate_layout(rust_cache, 1920.0, 1080.0, 2.0, 0.0, num_colors)
end = time.time()

duration_ms = (end - start) * 1000
print(f"Layout calculation (resize time): {duration_ms:.4f} ms")
print(f"Theoretical FPS: {1000/duration_ms:.2f} FPS")
print(f"FPS ceiling: {1000/duration_ms:.2f} FPS")