from __future__ import annotations

import unittest

from mypyc.common import (
    ENV_ATTR_NAME,
    GENERATOR_ATTRIBUTE_PREFIX,
    NEXT_LABEL_ATTR_NAME,
    TEMP_ATTR_NAME,
    generator_frame_attribute_prefix,
)
from mypyc.ir.rtypes import RInstance
from mypyc.test.testutil import build_ir_for_single_file2
from mypyc.transform.spill import insert_spills


class TestSpill(unittest.TestCase):
    def test_separate_generator_environment_keeps_spills_on_frame(self) -> None:
        # A nested generator needs a separate environment. Since make() is
        # evaluated before the yield, its result must be spilled across the
        # suspension point by the post-IRBuild spill pass.
        source = """\
def make() -> str:
    return "left"

def outer():
    def nested(value: str):
        result = ""
        for item in [value]:
            # The loop iterator is promoted from a register during IR building,
            # while the result of make() is spilled by the later transform.
            result += make() + (yield item)
        return result
    return nested("right")
"""
        module, _, _, _ = build_ir_for_single_file2(source.splitlines())
        frame = next(cl for cl in module.classes if cl.has_running_flag)
        assert frame.env_user_function is not None

        insert_spills(frame.env_user_function, frame)

        env_type = frame.attributes[ENV_ATTR_NAME]
        assert isinstance(env_type, RInstance)
        environment = env_type.class_ir

        # Private generator resume state lives on the generator frame, which
        # is protected by the running flag and can use plain attribute access.
        assert frame.attrs_are_thread_confined()
        assert NEXT_LABEL_ATTR_NAME in frame.attributes
        assert any(name.startswith(TEMP_ATTR_NAME + "2_") for name in frame.attributes)
        assert any(name.startswith(TEMP_ATTR_NAME + "3_") for name in frame.attributes)

        # Noncaptured source-level variables also stay on the private frame.
        frame_prefix = generator_frame_attribute_prefix(
            frame.fullname, is_final_class=frame.is_final_class
        )
        assert frame_prefix + "value" in frame.attributes
        assert GENERATOR_ATTRIBUTE_PREFIX + "value" not in environment.attributes

    def test_generator_locals_use_private_frame_unless_captured(self) -> None:
        source = """\
def outer():
    def nested(captured: str, private: str):
        def callback() -> str:
            return captured

        private_local = private
        yield callback()
        yield private_local
    return nested("captured", "private")
"""
        module, _, _, _ = build_ir_for_single_file2(source.splitlines())
        frame = next(cl for cl in module.classes if cl.has_running_flag)

        env_type = frame.attributes[ENV_ATTR_NAME]
        assert isinstance(env_type, RInstance)
        environment = env_type.class_ir
        frame_attrs = set(frame.attributes)
        environment_attrs = set(environment.attributes)

        assert frame.attrs_are_thread_confined()
        assert not environment.attrs_are_thread_confined()

        frame_prefix = generator_frame_attribute_prefix(
            frame.fullname, is_final_class=frame.is_final_class
        )
        assert GENERATOR_ATTRIBUTE_PREFIX + "captured" in environment_attrs
        assert frame_prefix + "captured" not in frame_attrs
        for name in ("private", "private_local", "callback"):
            assert frame_prefix + name in frame_attrs
            assert GENERATOR_ATTRIBUTE_PREFIX + name not in environment_attrs