Defunctionalization: normalize aggregate closure capture environments

[Ian Davis (idavis)]

Summary

The branch fixes bugs where aggregate closure captures could be hidden, reordered, read too late, reconstructed incorrectly, or inconsistently expanded across closure occurrences. It safely flattens eligible local tuples and UDTs into capture leaves while preserving snapshots, nominal identity, argument order, nested closures, functors, and cross-package correctness.

This branch fixes defunctionalization bugs involving closures that capture aggregates containing callable values, such as tuples, nested tuples, and local UDTs.

1. Callable values hidden inside aggregate captures

A callable nested inside a tuple or UDT could remain opaque to defunctionalization.

function ApplyCaptured() : Int {
    let environment = (10, x -> x + 1);
    let apply = x -> environment::Item2(x) + environment::Item1;
    apply(5)
}

The branch exposes eligible aggregate fields as separate capture parameters, allowing the nested callable to be analyzed and specialized.

2. Closure and lifted-target signature mismatches

Expanding an aggregate capture changes the lifted target’s input layout. If closure occurrences retain the old layout, arguments can be shifted or mistyped.

function Make(offset : Int) : Int -> Int {
    let environment = (offset, x -> 2 * x);
    x -> environment::Item2(x) + environment::Item1
}

operation Main() : Int {
    Make(1)(3) + Make(10)(3)
}

Every closure occurrence sharing the lifted target must use the same expanded capture layout.

3. Capture snapshots changing after reassignment

Aggregate fields must be captured when the closure is created, not read when it is later invoked.

operation Main() : Int {
    mutable environment = (1, x -> x + 10);
    let captured = x -> environment::Item2(x) + environment::Item1;

    set environment = (100, x -> x + 1000);

    captured(2) // Must return 13, not 1102.
}

The branch projects aggregate leaves into immutable locals at closure creation.

4. Reordered or repeated field evaluation

Normalizing an environment must not replay or reorder the expressions that produced its fields.

function TraceValue(label : String, value : Int) : Int {
    Message(label);
    value
}

operation Main() : Int {
    let environment = (
        TraceValue("first", 1),
        x -> x + TraceValue("captured", 2),
        TraceValue("last", 3)
    );

    let apply = x -> environment::Item2(x)
                         + environment::Item1
                         + environment::Item3;
    apply(4)
}

first, captured, and last must each be evaluated once and in source order.

5. Lost tuple shape or UDT identity

Flattening a captured UDT into leaves must not replace the UDT with an anonymous tuple inside the lifted body.

struct Environment {
    Offset : Int,
    Transform : Int -> Int
}

function Make(environment : Environment) : Int -> Int {
    x -> environment.Transform(x) + environment.Offset
}

operation Main() : Int {
    let environment = new Environment {
        Offset = 3,
        Transform = x -> x * 2
    };

    Make(environment)(4) // Must return 11.
}

The branch reconstructs local UDTs through their original constructors.

6. Incorrect flattening of arrays

Callable arrays are runtime containers and cannot be expanded into a fixed number of capture parameters.

operation Main() : Int {
    let functions = [
        x -> x + 1,
        x -> x + 10
    ];

    let apply = (index, value) -> functions[index](value);

    apply(0, 5) + apply(1, 5) // Must return 21.
}

The branch keeps arrays opaque, preserving indexing and each element’s distinct capture environment.

7. Broken nested closure recaptures

An inner closure may recapture an aggregate from an outer lifted closure. Expanding only the outer capture can leave the inner closure referring to a removed aggregate local.

function MakeOuter(offset : Int) : Int -> Int {
    let environment = (offset, x -> x * 2);

    let outer = x -> {
        let inner = y ->
            environment::Item2(y)
            + environment::Item1;

        inner(x)
    };

    outer
}

operation Main() : Int {
    MakeOuter(3)(4) // Must return 11.
}

The branch forwards the outer closure’s replacement leaf captures into the inner closure.

8. Incompatible normalization across shared targets

Multiple closure instances can share one lifted target while capturing different environments.

function Make(offset : Int) : Int -> Int {
    let environment = (offset, x -> x * 2);
    x -> environment::Item2(x) + environment::Item1
}

operation Main() : Int {
    let first = Make(1);
    let second = Make(10);

    first(3) + second(3) // Must return 23.
}

The branch updates every occurrence consistently or declines normalization when it cannot safely update the shared target.

9. Capture and ordinary-argument slot reordering

Expanded capture leaves must remain separate from the lambda’s ordinary parameters.

function Make(offset : Int, scale : Int) : (Int, Int) -> Int {
    let first = (offset, x -> x + 1);
    let second = (scale, x -> x * 2);

    (left, right) ->
        first::Item2(left) * first::Item1
        + second::Item2(right) * second::Item1
}

operation Main() : Int {
    Make(3, 5)(2, 4) // (2 + 1) * 3 + (4 * 2) * 5 = 49.
}

The branch preserves depth-first field order, unaffected captures, and ordinary argument positions.

10. Controlled and adjoint calls receiving corrupted environments

Captured operation values and their associated data must remain correctly positioned under functor application.

struct OperationEnvironment {
    Angle : Double,
    Apply : (Double, Qubit) => Unit is Adj + Ctl
}

function MakeRotation(environment : OperationEnvironment)
: Qubit => Unit is Adj + Ctl {
    q => environment.Apply(environment.Angle, q)
}

operation Main() : Result {
    use control = Qubit();
    use target = Qubit();

    X(control);

    let environment = new OperationEnvironment {
        Angle = PI(),
        Apply = Rz
    };

    let rotate = MakeRotation(environment);
    Controlled rotate([control], target);

    Reset(control);
    MResetZ(target)
}

The branch preserves the captured angle, operation value, ordinary target argument, and control register across body, adjoint, controlled, and controlled-adjoint forms.

11. Cross-package aggregate captures losing nominal ownership

A foreign UDT captured beside a callable must retain its original package-qualified identity.

// Library package
namespace Library {
    struct Payload {
        Value : Int
    }

    function MakePayload(value : Int) : Payload {
        new Payload { Value = value }
    }
}

// User package
namespace Application {
    import Library.*;

    operation Main() : Int {
        let environment = (
            MakePayload(7),
            x -> x + 1
        );

        let apply = x ->
            environment::Item2(x)
            + environment::Item1.Value;

        apply(3) // Must return 11.
    }
}

The branch keeps the foreign UDT opaque while normalizing the callable sibling.