`@transfunction` and `@superfunction` can read or mutate the wrong closure variables

[pomponchik]

When a template captures several variables from an enclosing function,
generated variants can connect a variable name to the wrong closure cell.
The generated code then silently reads or writes a different variable from the
one named in the source.

For example:

from asyncio import run

from transfunctions import (
    async_context,
    generator_context,
    sync_context,
    transfunction,
)


def make_template():
    async_value = "ASYNC"
    generator_value = "GENERATOR"
    sync_value = "SYNC"

    @transfunction
    def template():
        with sync_context:
            return sync_value
        with async_context:
            return async_value
        with generator_context:
            yield generator_value

    return template


template = make_template()

print(template.get_usual_function()())
print(run(template.get_async_function()()))
print(list(template.get_generator_function()()))

Expected:

SYNC
ASYNC
['GENERATOR']

Actual:

ASYNC
ASYNC
['ASYNC']

This is not limited to incorrect return values. A nonlocal assignment can
silently update an unrelated captured variable:

def make_writer():
    async_value = "ASYNC"
    sync_value = "SYNC"

    @transfunction
    def write():
        nonlocal sync_value

        with sync_context:
            sync_value = "CHANGED"
        with async_context:
            return async_value

    def state():
        return async_value, sync_value

    return write, state


write, state = make_writer()
write.get_usual_function()()

print(state())
# Actual:   ('CHANGED', 'SYNC')
# Expected: ('ASYNC', 'CHANGED')

The same problem affects @superfunction, because it uses
FunctionTransformer for the same context extraction and closure-rewriting
process. No exception or warning indicates that a different closure cell was
accessed.

Regression tests

The following tests currently fail. They should all pass after the bug is
fixed:

from asyncio import run

from transfunctions import (
    async_context,
    generator_context,
    superfunction,
    sync_context,
    transfunction,
)


def make_transfunction():
    async_value = "ASYNC"
    generator_value = "GENERATOR"
    sync_value = "SYNC"

    @transfunction
    def template():
        with sync_context:
            return sync_value
        with async_context:
            return async_value
        with generator_context:
            yield generator_value

    return template


def make_superfunction():
    async_value = "ASYNC"
    generator_value = "GENERATOR"
    sync_value = "SYNC"

    @superfunction
    def template():
        with sync_context:
            return sync_value
        with async_context:
            return async_value
        with generator_context:
            yield generator_value

    return template


def make_transfunction_writer():
    async_value = "ASYNC"
    sync_value = "SYNC"

    @transfunction
    def write():
        nonlocal sync_value

        with sync_context:
            sync_value = "CHANGED"
        with async_context:
            return async_value

    def state():
        return async_value, sync_value

    return write, state


def make_superfunction_writer():
    async_value = "ASYNC"
    sync_value = "SYNC"

    @superfunction
    def write():
        nonlocal sync_value

        with sync_context:
            sync_value = "CHANGED"
        with async_context:
            return async_value

    def state():
        return async_value, sync_value

    return write, state


def test_transfunction_variants_use_the_cells_named_in_the_source():
    template = make_transfunction()

    actual = (
        template.get_usual_function()(),
        run(template.get_async_function()()),
        list(template.get_generator_function()()),
    )

    assert actual == ("SYNC", "ASYNC", ["GENERATOR"])


def test_superfunction_variants_use_the_cells_named_in_the_source():
    template = make_superfunction()

    actual = (
        ~template(),
        run(template()),
        list(template()),
    )

    assert actual == ("SYNC", "ASYNC", ["GENERATOR"])


def test_transfunction_nonlocal_write_updates_the_named_cell():
    write, state = make_transfunction_writer()

    write.get_usual_function()()

    assert state() == ("ASYNC", "CHANGED")


def test_superfunction_nonlocal_write_updates_the_named_cell():
    write, state = make_superfunction_writer()

    ~write()

    assert state() == ("ASYNC", "CHANGED")

Current result:

4 failed

The read tests currently produce:

('ASYNC', 'ASYNC', ['ASYNC'])

Both write tests currently produce:

('CHANGED', 'SYNC')

Possible cause

Context extraction removes marker branches that do not belong to the requested
variant. This can reduce the set of free variables used by the newly compiled
function. Its LOAD_DEREF and STORE_DEREF bytecode operands therefore refer
to positions in that reduced, compacted free-variable list.

FunctionTransformer.rewrite_globals_and_closure() subsequently replaces the
generated code object's co_freevars with the complete free-variable list from
the original template and supplies the original closure cells in that order.
It does not rewrite the existing dereference operands to match the expanded
list. An operand that meant “the first remaining free variable” can therefore
refer to the first variable from the original template instead.

For example, after context extraction a generated function may have:

co_freevars = ('sync_value',)
LOAD_DEREF 0

After replacing only co_freevars, the effective mapping can become:

co_freevars = ('async_value', 'generator_value', 'sync_value')
LOAD_DEREF 0  # now reads async_value

The closure cells and bytecode indexes must retain the same name-to-cell
mapping after marker removal.

Environment

• transfunctions: 0.0.13
• Python: 3.10.8
• OS: macOS 15.3.1 (build 24D70)