represent the stack offset for a peek or poke at stackpointer0 deep and high are both empty note that that is an invalid stack reference below stack top only deep is nonempty above stack top only high is nonempty in complex cases both deep and high may be nonempty all this would be much simpler if all stack entries were the same size but with conditional and array effects they aren t the offsets are each represented by a list of stackeffect objects the name in the stackeffects is unused produce an index expression from the terms honoring pep 8 surrounding binary ops with spaces but not unary minus return e g stackpointer1 terms self offset asterms if self effect size terms insert0 stackpointer index makeindexterms if self effect size res index else res fstackpointerindex if not lax check that we re not reading or writing above stack top skip this for output variable initialization laxtrue assert self effect in self offset deep and not self offset high fpush or pop above current stack level res return res def asstackeffectself lax bool false stackeffect return stackeffect self asvariablelaxlax self effect type if self effect size else self effect cond self effect size dataclasses dataclass class copyitem src stackitem dst stackitem class effectmanager track offsets from stack pointer link to previous manager replace pushx popy with copyx y check that the sources and destinations are disjoint todo turn this into an error pass an analyzer instance see if we can get more copies of a earlier predecessor fix up patterns of copies through unused e g cpa unused cpunused b cpa b collect all variables skipping unused ones vars dictstr stackeffect def addeff stackeffect none if eff name unused if eff name in vars todo make this an error assert varseff name eff self instr name eff name varseff name eff else varseff name eff for copy in self copies addcopy src effect addcopy dst effect for peek in self peeks addpeek effect for poke in self pokes addpoke effect return vars def lessthana stackoffset b stackoffset bool todo handle more cases if a high b high return false return a deep lenb deep b deep def getmanagersparts listcomponent listeffectmanager managers listeffectmanager pred effectmanager none none for part in parts mgr effectmanagerpart instr part activecaches pred managers appendmgr pred mgr return managers def getstackeffectinfoformacromac macroinstruction tuplestr str parts part for part in mac parts if isinstancepart component managers getmanagersparts popped stackoffset for mgr in managers if lessthanmgr minoffset popped popped mgr minoffset clone compute pushed final popped pushed managers1 finaloffset clone for effect in popped deep pushed highereffect for effect in popped high pushed deepereffect return popped negate asindex pushed asindex def writesingleinstr instr instruction out formatter tier tiers tierone none try writecomponents componentinstr instr activecaches out tier 0 instr family except assertionerror as err raise assertionerrorferror writing instruction instr name from err def writemacroinstrmac macroinstruction out formatter none parts part for part in mac parts if isinstancepart component and part instr name setip out emit with out blockftargetmac name needsthis anypart instr needsthisinstr for part in parts if needsthis and not mac predicted out emitfpycodeunit thisinstr frameinstrptr nextinstr else out emitfframeinstrptr nextinstr out emitfnextinstr mac cacheoffset1 out emitfinstructionstatsmac name if mac predicted out emitfpredictedmac name if needsthis out emitfpycodeunit thisinstr nextinstr mac cacheoffset1 out staticassertfamilysizemac name mac family mac cacheoffset try nextinstrisset writecomponents parts out tierone mac cacheoffset mac family except assertionerror as err raise assertionerrorferror writing macro mac name from err if not parts1 instr alwaysexits if parts1 instr checkevalbreaker out emitcheckevalbreaker out emitdispatch def writecomponents parts listcomponent out formatter tier tiers cacheoffset int family family none bool managers getmanagersparts allvars dictstr stackeffect for mgr in managers for name eff in mgr collectvars items if name in allvars todo turn this into an error variable conflict assert allvarsname eff name mgr instr name allvarsname eff else allvarsname eff declare all variables for name eff in allvars items out declareeff none nextinstrisset false for mgr in managers if lenparts 1 out emitf mgr instr name for copy in mgr copies copysrceffect copy src effect if copysrceffect name copy dst effect name if copysrceffect name unused copysrceffect copy src asstackeffect out assigncopy dst effect copysrceffect for peek in mgr peeks out assign peek effect peek asstackeffect initialize array outputs for poke in mgr pokes if poke effect size and poke effect name not in mgr instr unmovednames out assign poke effect poke asstackeffectlaxtrue if mgr instr name in pushframe popframe adjust stack to minoffset this means that all input effects of this instruction are materialized but not its output effects that s as intended since these two are so special out stackadjustmgr minoffset deep mgr minoffset high however for tier 2 pretend the stack is at final offset mgr adjustinversemgr finaloffset if tier tierone todo check in analyzer that push popframe is last assert mgr is managers1 fexpected mgr instr name r to be the last uop assertnopokesmanagers if mgr instr name savereturnoffset nextinstrisset true if tier tierone assertnopokesmanagers if lenparts 1 mgr instr writebodyout 0 mgr activecaches tier family else with out block mgr instr writebodyout 4 mgr activecaches tier family if mgr is managers1 and not nextinstrisset and not mgr instr alwaysexits adjust the stack to its final depth then write the pokes for all preceding uops note that for array output effects we may still write past the stack top out stackadjustmgr finaloffset deep mgr finaloffset high writeallpokesmgr finaloffset managers out return nextinstrisset def assertnopokesmanagers listeffectmanager none for mgr in managers for poke in mgr pokes if not poke effect size and poke effect name not in mgr instr unmovednames assert poke effect name unused funexpected poke of poke effect name in mgr instr name r def writeallpokes offset stackoffset managers listeffectmanager out formatter none emit all remaining pushes pokes for m in managers m adjustinverseoffset writepokesm out def writepokesmgr effectmanager out formatter none for poke in mgr pokes if not poke effect size and poke effect name not in mgr instr unmovednames out assign poke asstackeffect poke effect def writesingleinstrforabstractinterpinstr instruction out formatter none try writecomponentsforabstractinterp componentinstr instr activecaches out except assertionerror as err raise assertionerror ferror writing abstract instruction instr name from err def writecomponentsforabstractinterp parts listcomponent out formatter none managers getmanagersparts for mgr in managers if mgr is managers1 out stackadjustmgr finaloffset deep mgr finaloffset high mgr adjustinversemgr finaloffset null out the output stack effects for poke in mgr pokes if not poke effect size and poke effect name not in mgr instr unmovednames out emit fpartitionnodeoverwritepypartitionnodet fpartitionnodenullroot peekpoke offset asindex true represent the stack offset for a peek or poke at stack_pointer 0 deep and high are both empty note that that is an invalid stack reference below stack top only deep is non empty above stack top only high is non empty in complex cases both deep and high may be non empty all this would be much simpler if all stack entries were the same size but with conditional and array effects they aren t the offsets are each represented by a list of stackeffect objects the name in the stackeffects is unused produce an index expression from the terms honoring pep 8 surrounding binary ops with spaces but not unary minus return e g stack_pointer 1 check that we re not reading or writing above stack top skip this for output variable initialization lax true manage stack effects and offsets for an instruction see merge track offsets from stack pointer link to previous manager replace push x pop y with copy x y check that the sources and destinations are disjoint todo turn this into an error pass an analyzer instance see if we can get more copies of a earlier predecessor break fix up patterns of copies through unused e g cp a unused cp unused b cp a b collect all variables skipping unused ones todo make this an error todo handle more cases get the stack effect info for a macro instruction returns a tuple popped pushed where each is a string giving a symbolic expression for the number of values popped pushed compute pushed final popped todo turn this into an error variable conflict declare all variables initialize array outputs adjust stack to min_offset this means that all input effects of this instruction are materialized but not its output effects that s as intended since these two are so special however for tier 2 pretend the stack is at final offset todo check in analyzer that _ push pop _frame is last adjust the stack to its final depth then write the pokes for all preceding uops note that for array output effects we may still write past the stack top emit all remaining pushes pokes null out the output stack effects
import dataclasses
import typing

from flags import variable_used_unspecialized
from formatting import (
    Formatter,
    UNUSED,
    maybe_parenthesize,
    parenthesize_cond,
)
from instructions import (
    ActiveCacheEffect,
    Instruction,
    MacroInstruction,
    Component,
    Tiers,
    TIER_ONE,
)
from parsing import StackEffect, CacheEffect, Family


@dataclasses.dataclass
class StackOffset:
    

    deep: list[StackEffect] = dataclasses.field(default_factory=list)
    high: list[StackEffect] = dataclasses.field(default_factory=list)

    def clone(self) -> "StackOffset":
        return StackOffset(list(self.deep), list(self.high))

    def negate(self) -> "StackOffset":
        return StackOffset(list(self.high), list(self.deep))

    def deeper(self, eff: StackEffect) -> None:
        if eff in self.high:
            self.high.remove(eff)
        else:
            self.deep.append(eff)

    def higher(self, eff: StackEffect) -> None:
        if eff in self.deep:
            self.deep.remove(eff)
        else:
            self.high.append(eff)

    def as_terms(self) -> list[tuple[str, str]]:
        num = 0
        terms: list[tuple[str, str]] = []
        for eff in self.deep:
            if eff.size:
                terms.append(("-", maybe_parenthesize(eff.size)))
            elif eff.cond and eff.cond not in ("0", "1"):
                terms.append(("-", f"({parenthesize_cond(eff.cond)} ? 1 : 0)"))
            elif eff.cond != "0":
                num -= 1
        for eff in self.high:
            if eff.size:
                terms.append(("+", maybe_parenthesize(eff.size)))
            elif eff.cond and eff.cond not in ("0", "1"):
                terms.append(("+", f"({parenthesize_cond(eff.cond)} ? 1 : 0)"))
            elif eff.cond != "0":
                num += 1
        if num < 0:
            terms.insert(0, ("-", str(-num)))
        elif num > 0:
            terms.append(("+", str(num)))
        return terms

    def as_index(self) -> str:
        terms = self.as_terms()
        return make_index(terms)

    def equivalent_to(self, other: "StackOffset") -> bool:
        if self.deep == other.deep and self.high == other.high:
            return True
        deep = list(self.deep)
        for x in other.deep:
            try:
                deep.remove(x)
            except ValueError:
                return False
        if deep:
            return False
        high = list(self.high)
        for x in other.high:
            try:
                high.remove(x)
            except ValueError:
                return False
        if high:
            return False
        return True


def make_index(terms: list[tuple[str, str]]) -> str:
    
    
    index = ""
    for sign, term in terms:
        if index:
            index += f" {sign} {term}"
        elif sign == "+":
            index = term
        else:
            index = sign + term
    return index or "0"


@dataclasses.dataclass
class StackItem:
    offset: StackOffset
    effect: StackEffect

    def as_variable(self, lax: bool = False) -> str:
        
        terms = self.offset.as_terms()
        if self.effect.size:
            terms.insert(0, ("+", "stack_pointer"))
        index = make_index(terms)
        if self.effect.size:
            res = index
        else:
            res = f"stack_pointer[{index}]"
        if not lax:
            
            
            assert (
                self.effect in self.offset.deep and not self.offset.high
            ), f"Push or pop above current stack level: {res}"
        return res

    def as_stack_effect(self, lax: bool = False) -> StackEffect:
        return StackEffect(
            self.as_variable(lax=lax),
            self.effect.type if self.effect.size else "",
            self.effect.cond,
            self.effect.size,
        )


@dataclasses.dataclass
class CopyItem:
    src: StackItem
    dst: StackItem


class EffectManager:
    

    instr: Instruction
    active_caches: list[ActiveCacheEffect]
    peeks: list[StackItem]
    pokes: list[StackItem]
    copies: list[CopyItem]  
    
    min_offset: StackOffset
    final_offset: StackOffset
    
    pred: "EffectManager | None" = None

    def __init__(
        self,
        instr: Instruction,
        active_caches: list[ActiveCacheEffect],
        pred: "EffectManager | None" = None,
    ):
        self.instr = instr
        self.active_caches = active_caches
        self.peeks = []
        self.pokes = []
        self.copies = []
        self.final_offset = pred.final_offset.clone() if pred else StackOffset()
        for eff in reversed(instr.input_effects):
            self.final_offset.deeper(eff)
            self.peeks.append(StackItem(offset=self.final_offset.clone(), effect=eff))
        self.min_offset = self.final_offset.clone()
        for eff in instr.output_effects:
            self.pokes.append(StackItem(offset=self.final_offset.clone(), effect=eff))
            self.final_offset.higher(eff)

        self.pred = pred
        while pred:
            
            
            sources: set[str] = set()
            destinations: set[str] = set()
            while (
                pred.pokes
                and self.peeks
                and pred.pokes[-1].effect == self.peeks[0].effect
            ):
                src = pred.pokes.pop(-1)
                dst = self.peeks.pop(0)
                assert src.offset.equivalent_to(dst.offset), (src, dst)
                pred.final_offset.deeper(src.effect)
                if dst.effect.name != src.effect.name:
                    if dst.effect.name != UNUSED:
                        destinations.add(dst.effect.name)
                    if src.effect.name != UNUSED:
                        sources.add(src.effect.name)
                self.copies.append(CopyItem(src, dst))
            
            assert sources & destinations == set(), (
                pred.instr.name,
                self.instr.name,
                sources,
                destinations,
            )
            
            if self.peeks and not pred.pokes and not pred.peeks:
                pred = pred.pred
            else:
                pred = None  

        
        
        if any(copy.src.effect.name == UNUSED for copy in self.copies):
            pred = self.pred
            while pred is not None:
                for copy in self.copies:
                    if copy.src.effect.name == UNUSED:
                        for pred_copy in pred.copies:
                            if pred_copy.dst == copy.src:
                                copy.src = pred_copy.src
                                break
                pred = pred.pred

    def adjust_deeper(self, eff: StackEffect) -> None:
        for peek in self.peeks:
            peek.offset.deeper(eff)
        for poke in self.pokes:
            poke.offset.deeper(eff)
        for copy in self.copies:
            copy.src.offset.deeper(eff)
            copy.dst.offset.deeper(eff)
        self.min_offset.deeper(eff)
        self.final_offset.deeper(eff)

    def adjust_higher(self, eff: StackEffect) -> None:
        for peek in self.peeks:
            peek.offset.higher(eff)
        for poke in self.pokes:
            poke.offset.higher(eff)
        for copy in self.copies:
            copy.src.offset.higher(eff)
            copy.dst.offset.higher(eff)
        self.min_offset.higher(eff)
        self.final_offset.higher(eff)

    def adjust(self, offset: StackOffset) -> None:
        deep = list(offset.deep)
        high = list(offset.high)
        for down in deep:
            self.adjust_deeper(down)
        for up in high:
            self.adjust_higher(up)

    def adjust_inverse(self, offset: StackOffset) -> None:
        deep = list(offset.deep)
        high = list(offset.high)
        for down in deep:
            self.adjust_higher(down)
        for up in high:
            self.adjust_deeper(up)

    def collect_vars(self) -> dict[str, StackEffect]:
        
        vars: dict[str, StackEffect] = {}

        def add(eff: StackEffect) -> None:
            if eff.name != UNUSED:
                if eff.name in vars:
                    
                    assert vars[eff.name] == eff, (
                        self.instr.name,
                        eff.name,
                        vars[eff.name],
                        eff,
                    )
                else:
                    vars[eff.name] = eff

        for copy in self.copies:
            add(copy.src.effect)
            add(copy.dst.effect)
        for peek in self.peeks:
            add(peek.effect)
        for poke in self.pokes:
            add(poke.effect)

        return vars


def less_than(a: StackOffset, b: StackOffset) -> bool:
    
    if a.high != b.high:
        return False
    return a.deep[: len(b.deep)] == b.deep


def get_managers(parts: list[Component]) -> list[EffectManager]:
    managers: list[EffectManager] = []
    pred: EffectManager | None = None
    for part in parts:
        mgr = EffectManager(part.instr, part.active_caches, pred)
        managers.append(mgr)
        pred = mgr
    return managers


def get_stack_effect_info_for_macro(mac: MacroInstruction) -> tuple[str, str]:
    
    parts = [part for part in mac.parts if isinstance(part, Component)]
    managers = get_managers(parts)
    popped = StackOffset()
    for mgr in managers:
        if less_than(mgr.min_offset, popped):
            popped = mgr.min_offset.clone()
    
    pushed = managers[-1].final_offset.clone()
    for effect in popped.deep:
        pushed.higher(effect)
    for effect in popped.high:
        pushed.deeper(effect)
    return popped.negate().as_index(), pushed.as_index()


def write_single_instr(
    instr: Instruction, out: Formatter, tier: Tiers = TIER_ONE
) -> None:
    try:
        write_components(
            [Component(instr, instr.active_caches)],
            out,
            tier,
            0,
            instr.family,
        )
    except AssertionError as err:
        raise AssertionError(f"Error writing instruction {instr.name}") from err


def write_macro_instr(mac: MacroInstruction, out: Formatter) -> None:
    parts = [
        part
        for part in mac.parts
        if isinstance(part, Component) and part.instr.name != "_SET_IP"
    ]
    out.emit("")
    with out.block(f"TARGET({mac.name})"):
        needs_this = any(part.instr.needs_this_instr for part in parts)
        if needs_this and not mac.predicted:
            out.emit(f"_Py_CODEUNIT *this_instr = frame->instr_ptr = next_instr;")
        else:
            out.emit(f"frame->instr_ptr = next_instr;")
        out.emit(f"next_instr += {mac.cache_offset+1};")
        out.emit(f"INSTRUCTION_STATS({mac.name});")
        if mac.predicted:
            out.emit(f"PREDICTED({mac.name});")
            if needs_this:
                out.emit(f"_Py_CODEUNIT *this_instr = next_instr - {mac.cache_offset+1};")
        out.static_assert_family_size(mac.name, mac.family, mac.cache_offset)
        try:
            next_instr_is_set = write_components(
                parts, out, TIER_ONE, mac.cache_offset, mac.family
            )
        except AssertionError as err:
            raise AssertionError(f"Error writing macro {mac.name}") from err
        if not parts[-1].instr.always_exits:
            if parts[-1].instr.check_eval_breaker:
                out.emit("CHECK_EVAL_BREAKER();")
            out.emit("DISPATCH();")


def write_components(
    parts: list[Component],
    out: Formatter,
    tier: Tiers,
    cache_offset: int,
    family: Family | None,
) -> bool:
    managers = get_managers(parts)

    all_vars: dict[str, StackEffect] = {}
    for mgr in managers:
        for name, eff in mgr.collect_vars().items():
            if name in all_vars:
                
                assert all_vars[name] == eff, (
                    name,
                    mgr.instr.name,
                    all_vars[name],
                    eff,
                )
            else:
                all_vars[name] = eff

    
    for name, eff in all_vars.items():
        out.declare(eff, None)

    next_instr_is_set = False
    for mgr in managers:
        if len(parts) > 1:
            out.emit(f"// {mgr.instr.name}")

        for copy in mgr.copies:
            copy_src_effect = copy.src.effect
            if copy_src_effect.name != copy.dst.effect.name:
                if copy_src_effect.name == UNUSED:
                    copy_src_effect = copy.src.as_stack_effect()
                out.assign(copy.dst.effect, copy_src_effect)
        for peek in mgr.peeks:
            out.assign(
                peek.effect,
                peek.as_stack_effect(),
            )
        
        for poke in mgr.pokes:
            if poke.effect.size and poke.effect.name not in mgr.instr.unmoved_names:
                out.assign(
                    poke.effect,
                    poke.as_stack_effect(lax=True),
                )

        if mgr.instr.name in ("_PUSH_FRAME", "_POP_FRAME"):
            
            
            
            
            out.stack_adjust(mgr.min_offset.deep, mgr.min_offset.high)
            
            mgr.adjust_inverse(mgr.final_offset)
            if tier == TIER_ONE:
                
                assert (
                    mgr is managers[-1]
                ), f"Expected {mgr.instr.name!r} to be the last uop"
                assert_no_pokes(managers)

        if mgr.instr.name == "_SAVE_RETURN_OFFSET":
            next_instr_is_set = True
            if tier == TIER_ONE:
                assert_no_pokes(managers)

        if len(parts) == 1:
            mgr.instr.write_body(out, 0, mgr.active_caches, tier, family)
        else:
            with out.block(""):
                mgr.instr.write_body(out, -4, mgr.active_caches, tier, family)

        if mgr is managers[-1] and not next_instr_is_set and not mgr.instr.always_exits:
            
            
            
            
            out.stack_adjust(mgr.final_offset.deep, mgr.final_offset.high)
            write_all_pokes(mgr.final_offset, managers, out)

    return next_instr_is_set


def assert_no_pokes(managers: list[EffectManager]) -> None:
    for mgr in managers:
        for poke in mgr.pokes:
            if not poke.effect.size and poke.effect.name not in mgr.instr.unmoved_names:
                assert (
                    poke.effect.name == UNUSED
                ), f"Unexpected poke of {poke.effect.name} in {mgr.instr.name!r}"


def write_all_pokes(
    offset: StackOffset, managers: list[EffectManager], out: Formatter
) -> None:
    
    for m in managers:
        m.adjust_inverse(offset)
        write_pokes(m, out)


def write_pokes(mgr: EffectManager, out: Formatter) -> None:
    for poke in mgr.pokes:
        if not poke.effect.size and poke.effect.name not in mgr.instr.unmoved_names:
            out.assign(
                poke.as_stack_effect(),
                poke.effect,
            )


def write_single_instr_for_abstract_interp(instr: Instruction, out: Formatter) -> None:
    try:
        _write_components_for_abstract_interp(
            [Component(instr, instr.active_caches)],
            out,
        )
    except AssertionError as err:
        raise AssertionError(
            f"Error writing abstract instruction {instr.name}"
        ) from err


def _write_components_for_abstract_interp(
    parts: list[Component],
    out: Formatter,
) -> None:
    managers = get_managers(parts)
    for mgr in managers:
        if mgr is managers[-1]:
            out.stack_adjust(mgr.final_offset.deep, mgr.final_offset.high)
            mgr.adjust_inverse(mgr.final_offset)
        
        for poke in mgr.pokes:
            if not poke.effect.size and poke.effect.name not in mgr.instr.unmoved_names:
                out.emit(
                    f"PARTITIONNODE_OVERWRITE((_Py_PARTITIONNODE_t *)"
                    f"PARTITIONNODE_NULLROOT, PEEK(-({poke.offset.as_index()})), true);"
                )