Add Closure Variable Handling in Function Flattening

[jeremymanning]

Overview

The function flattening system has a TODO marker indicating missing support for closure variable arguments.

Current Status

Location: clustrix/function_flattening.py:751

Current Code:

# Need to add closure variables as arguments
node.func.id = self.hoisted_mapping[func_name]
# TODO: Add closure variable arguments

Problem Description

The function flattening process is not handling closure variables properly. When nested functions are hoisted to module level, they may lose access to variables from their original enclosing scope, leading to runtime errors.

Impact

• Nested functions that depend on closure variables will fail during remote execution
• NameError exceptions when accessing variables that were in the original closure
• Incomplete function transformation that breaks semantic correctness

Example Problem Case

@cluster(cores=2)
def outer_function(data):
    threshold = 0.5  # Closure variable
    
    def inner_function(item):
        return item > threshold  # References closure variable
    
    return [inner_function(x) for x in data]

When inner_function is hoisted, it loses access to threshold.

Implementation Requirements

1. Closure Variable Detection

• Identify variables referenced in nested functions that come from enclosing scopes
• Distinguish between closure variables, global variables, and local variables
• Handle multiple levels of nested closures

2. Function Signature Modification

• Add closure variables as explicit parameters to hoisted functions
• Update function call sites to pass closure variables as arguments
• Preserve original function semantics

3. Call Site Updates

• Modify all call sites of hoisted functions
• Pass closure variable values as additional arguments
• Handle complex closure variable expressions

Technical Implementation

Closure Variable Analysis

def detect_closure_variables(func_node, enclosing_scope):
    """Detect variables used in func_node that come from enclosing_scope."""
    # Analyze variable references
    # Return mapping of variable names to values

Function Signature Transformation

def add_closure_parameters(func_node, closure_vars):
    """Add closure variables as parameters to function definition."""
    # Modify function signature
    # Add parameters for each closure variable

Call Site Updates

def update_call_sites(call_node, func_name, closure_vars):
    """Update function calls to pass closure variables."""
    # Add closure variables as function arguments
    # Preserve original argument order and semantics

Testing Requirements

• Test simple closure variable cases
• Test multiple levels of nested closures
• Test complex closure variable expressions
• Test closure variables with different types
• Test interaction with existing function flattening
• Test performance impact of closure variable handling

Related Components

• clustrix/function_flattening.py - Main implementation location
• clustrix/loop_analysis.py - AST analysis utilities
• clustrix/dependency_resolution.py - Related to variable dependency tracking

Acceptance Criteria

• TODO comment is resolved with actual implementation
• Closure variables are correctly detected in nested functions
• Hoisted functions receive closure variables as parameters
• All call sites are updated to pass closure variables
• Original function semantics are preserved
• All existing function flattening tests continue to pass
• New tests cover closure variable scenarios

Priority: High (blocks correct nested function handling)
Complexity: High

Source: Content found during codebase TODO/FIXME scan in Issue #74 content preservation automation.