parser for bytecodes inst from dataclasses import dataclass field from typing import namedtuple callable typevar literal cast import lexer as lx from plexer import plexer p typevarp boundparser n typevarn boundnode def contextualfunc callablep n none callablep n none decorator to wrap grammar methods resets position if func returns none def contextualwrapperself p n none begin self getpos res funcself if res is none self setposbegin return none end self getpos res context contextbegin end self return res return contextualwrapper class contextnamedtuple begin int end int owner plexer def reprself str return fself owner filename self beginself end dataclass class node context context none fieldinitfalse comparefalse defaultnone property def textself str return self totext def totextself dedent int 0 str context self context if not context return return lx totextself tokens dedent property def tokensself listlx token context self context if not context return tokens context owner tokens begin context begin end context end return tokensbegin end dataclass class blocknode this just holds a context which has the list of tokens pass dataclass class stackeffectnode name str fieldcomparefalse eq only uses type cond size type str optional type cond str optional if cond size str optional size note size cannot be combined with type or cond def reprself str items self name self type self cond self size while items and items1 del items1 return fstackeffect joinrepritem for item in items dataclass class expressionnode size str dataclass class cacheeffectnode name str size int dataclass class opnamenode name str inputeffect stackeffect cacheeffect outputeffect stackeffect uop opname cacheeffect dataclass class instheadernode annotations liststr kind literalinst op name str inputs listinputeffect outputs listoutputeffect dataclass class instdefnode annotations liststr kind literalinst op name str inputs listinputeffect outputs listoutputeffect block block dataclass class macronode name str uops listuop dataclass class familynode name str size str variable giving the cache size in code units members liststr dataclass class pseudonode name str flags liststr instr flags to set on the pseudo instruction targets liststr opcodes this can be replaced by class parserplexer contextual def definitionself instdef macro pseudo family none if macro self macrodef return macro if family self familydef return family if pseudo self pseudodef return pseudo if inst self instdef return inst return none contextual def instdefself instdef none if hdr self instheader if block self block return instdef hdr annotations hdr kind hdr name hdr inputs hdr outputs block raise self makesyntaxerrorexpected block return none contextual def instheaderself instheader none annotation instname inputs outputs annotation opname inputs outputs annotations while anno self expectlx annotation annotations appendanno text tkn self expectlx inst if not tkn tkn self expectlx op if tkn kind castliteralinst op tkn text if self expectlx lparen and tkn self expectlx identifier name tkn text if self expectlx comma inp outp self ioeffect if self expectlx rparen if tkn self peek and tkn kind lx lbrace return instheaderannotations kind name inp outp return none def ioeffectself tuplelistinputeffect listoutputeffect inputs outputs if self expectlx lparen inputs self inputs or if self expectlx minusminus outputs self outputs or if self expectlx rparen return inputs outputs raise self makesyntaxerrorexpected stack effect def inputsself listinputeffect none input input here self getpos if inp self input inp castinputeffect inp near self getpos if self expectlx comma if rest self inputs return inp rest self setposnear return inp self setposhere return none contextual def inputself inputeffect none return self cacheeffect or self stackeffect def outputsself listoutputeffect none output output here self getpos if outp self output near self getpos if self expectlx comma if rest self outputs return outp rest self setposnear return outp self setposhere return none contextual def outputself outputeffect none return self stackeffect contextual def cacheeffectself cacheeffect none identifier number if tkn self expectlx identifier if self expectlx divide num self requirelx number text try size intnum except valueerror raise self makesyntaxerrorfexpected integer got num r else return cacheeffecttkn text size return none contextual def stackeffectself stackeffect none identifier identifier times if expression identifier expression if tkn self expectlx identifier typetext if self expectlx colon typetext self requirelx identifier text strip if self expectlx times typetext condtext if self expectlx if self requirelx lparen if not cond self expression raise self makesyntaxerrorexpected condition self requirelx rparen condtext cond text strip sizetext if self expectlx lbracket if typetext or condtext raise self makesyntaxerrorunexpected if not size self expression raise self makesyntaxerrorexpected expression self requirelx rbracket typetext pyobject sizetext size text strip return stackeffecttkn text typetext condtext sizetext return none contextual def expressionself expression none tokens listlx token level 1 while tkn self peek if tkn kind in lx lbracket lx lparen level 1 elif tkn kind in lx rbracket lx rparen level 1 if level 0 break tokens appendtkn self next if not tokens return none return expressionlx totexttokens strip def opsself listopname none if op self op ops op while self expectlx plus if op self op ops appendop return ops contextual def opself opname none if tkn self expectlx identifier return opnametkn text return none contextual def macrodefself macro none if tkn self expectlx macro if self expectlx lparen if tkn self expectlx identifier if self expectlx rparen if self expectlx equals if uops self uops self requirelx semi res macrotkn text uops return res return none def uopsself listuop none if uop self uop uop castuop uop uops uop while self expectlx plus if uop self uop uop castuop uop uops appenduop else raise self makesyntaxerrorexpected op name or cache effect return uops return none contextual def uopself uop none if tkn self expectlx identifier if self expectlx divide if num self expectlx number try size intnum text except valueerror raise self makesyntaxerror fexpected integer got num text r else return cacheeffecttkn text size raise self makesyntaxerrorexpected integer else return opnametkn text return none contextual def familydefself family none if tkn self expectlx identifier and tkn text family size none if self expectlx lparen if tkn self expectlx identifier if self expectlx comma if not size self expectlx identifier if not size self expectlx number raise self makesyntaxerrorexpected identifier or number if self expectlx rparen if self expectlx equals if not self expectlx lbrace raise self makesyntaxerrorexpected if members self members if self expectlx rbrace and self expectlx semi return family tkn text size text if size else members return none def flagsself liststr here self getpos if self expectlx lparen if tkn self expectlx identifier flags tkn text while self expectlx comma if tkn self expectlx identifier flags appendtkn text else break if not self expectlx rparen raise self makesyntaxerrorexpected comma or right paren return flags self setposhere return contextual def pseudodefself pseudo none if tkn self expectlx identifier and tkn text pseudo size none if self expectlx lparen if tkn self expectlx identifier if self expectlx comma flags self flags else flags if self expectlx rparen if self expectlx equals if not self expectlx lbrace raise self makesyntaxerrorexpected if members self members if self expectlx rbrace and self expectlx semi return pseudotkn text flags members return none def membersself liststr none here self getpos if tkn self expectlx identifier members tkn text while self expectlx comma if tkn self expectlx identifier members appendtkn text else break peek self peek if not peek or peek kind lx rbrace raise self makesyntaxerrorexpected comma or right paren return members self setposhere return none contextual def blockself block none if self cblob return block return none def cblobself listlx token tokens listlx token level 0 while tkn self nextrawtrue tokens appendtkn if tkn kind in lx lbrace lx lparen lx lbracket level 1 elif tkn kind in lx rbrace lx rparen lx rbracket level 1 if level 0 break return tokens if name main import sys if sys argv1 filename sys argv1 if filename c and sys argv2 src sys argv2 filename string else with openfilename r as f src f read srclines src splitlines begin srclines index begin bytecodes end srclines index end bytecodes src n joinsrclinesbegin 1 end else filename default src if x x foo commentn parser parsersrc filename x parser definition printx decorator to wrap grammar methods resets position if func returns none this just holds a context which has the list of tokens __eq__ only uses type cond size optional type optional if cond optional size note size cannot be combined with type or cond variable giving the cache size in code units instr flags to set on the pseudo instruction opcodes this can be replaced by annotation inst name inputs outputs annotation op name inputs outputs inputs outputs input input output output identifier number identifier identifier times if expression identifier expression def ops self list opname none if op self op ops op while self expect lx plus if op self op ops append op return ops
from dataclasses import dataclass, field
from typing import NamedTuple, Callable, TypeVar, Literal, cast

import lexer as lx
from plexer import PLexer


P = TypeVar("P", bound="Parser")
N = TypeVar("N", bound="Node")


def contextual(func: Callable[[P], N | None]) -> Callable[[P], N | None]:
    
    
    def contextual_wrapper(self: P) -> N | None:
        begin = self.getpos()
        res = func(self)
        if res is None:
            self.setpos(begin)
            return None
        end = self.getpos()
        res.context = Context(begin, end, self)
        return res

    return contextual_wrapper


class Context(NamedTuple):
    begin: int
    end: int
    owner: PLexer

    def __repr__(self) -> str:
        return f"<{self.owner.filename}: {self.begin}-{self.end}>"


@dataclass
class Node:
    context: Context | None = field(init=False, compare=False, default=None)

    @property
    def text(self) -> str:
        return self.to_text()

    def to_text(self, dedent: int = 0) -> str:
        context = self.context
        if not context:
            return ""
        return lx.to_text(self.tokens, dedent)

    @property
    def tokens(self) -> list[lx.Token]:
        context = self.context
        if not context:
            return []
        tokens = context.owner.tokens
        begin = context.begin
        end = context.end
        return tokens[begin:end]


@dataclass
class Block(Node):
    
    pass


@dataclass
class StackEffect(Node):
    name: str = field(compare=False)  
    type: str = ""  
    cond: str = ""  
    size: str = ""  
    

    def __repr__(self) -> str:
        items = [self.name, self.type, self.cond, self.size]
        while items and items[-1] == "":
            del items[-1]
        return f"StackEffect({', '.join(repr(item) for item in items)})"


@dataclass
class Expression(Node):
    size: str


@dataclass
class CacheEffect(Node):
    name: str
    size: int


@dataclass
class OpName(Node):
    name: str


InputEffect = StackEffect | CacheEffect
OutputEffect = StackEffect
UOp = OpName | CacheEffect


@dataclass
class InstHeader(Node):
    annotations: list[str]
    kind: Literal["inst", "op"]
    name: str
    inputs: list[InputEffect]
    outputs: list[OutputEffect]


@dataclass
class InstDef(Node):
    annotations: list[str]
    kind: Literal["inst", "op"]
    name: str
    inputs: list[InputEffect]
    outputs: list[OutputEffect]
    block: Block


@dataclass
class Macro(Node):
    name: str
    uops: list[UOp]


@dataclass
class Family(Node):
    name: str
    size: str  
    members: list[str]


@dataclass
class Pseudo(Node):
    name: str
    flags: list[str]   
    targets: list[str] 


class Parser(PLexer):
    @contextual
    def definition(self) -> InstDef | Macro | Pseudo | Family | None:
        if macro := self.macro_def():
            return macro
        if family := self.family_def():
            return family
        if pseudo := self.pseudo_def():
            return pseudo
        if inst := self.inst_def():
            return inst
        return None

    @contextual
    def inst_def(self) -> InstDef | None:
        if hdr := self.inst_header():
            if block := self.block():
                return InstDef(
                    hdr.annotations,
                    hdr.kind,
                    hdr.name,
                    hdr.inputs,
                    hdr.outputs,
                    block,
                )
            raise self.make_syntax_error("Expected block")
        return None

    @contextual
    def inst_header(self) -> InstHeader | None:
        
        
        annotations = []
        while anno := self.expect(lx.ANNOTATION):
            annotations.append(anno.text)
        tkn = self.expect(lx.INST)
        if not tkn:
            tkn = self.expect(lx.OP)
        if tkn:
            kind = cast(Literal["inst", "op"], tkn.text)
            if self.expect(lx.LPAREN) and (tkn := self.expect(lx.IDENTIFIER)):
                name = tkn.text
                if self.expect(lx.COMMA):
                    inp, outp = self.io_effect()
                    if self.expect(lx.RPAREN):
                        if (tkn := self.peek()) and tkn.kind == lx.LBRACE:
                            return InstHeader(annotations, kind, name, inp, outp)
        return None

    def io_effect(self) -> tuple[list[InputEffect], list[OutputEffect]]:
        
        if self.expect(lx.LPAREN):
            inputs = self.inputs() or []
            if self.expect(lx.MINUSMINUS):
                outputs = self.outputs() or []
                if self.expect(lx.RPAREN):
                    return inputs, outputs
        raise self.make_syntax_error("Expected stack effect")

    def inputs(self) -> list[InputEffect] | None:
        
        here = self.getpos()
        if inp := self.input():
            inp = cast(InputEffect, inp)
            near = self.getpos()
            if self.expect(lx.COMMA):
                if rest := self.inputs():
                    return [inp] + rest
            self.setpos(near)
            return [inp]
        self.setpos(here)
        return None

    @contextual
    def input(self) -> InputEffect | None:
        return self.cache_effect() or self.stack_effect()

    def outputs(self) -> list[OutputEffect] | None:
        
        here = self.getpos()
        if outp := self.output():
            near = self.getpos()
            if self.expect(lx.COMMA):
                if rest := self.outputs():
                    return [outp] + rest
            self.setpos(near)
            return [outp]
        self.setpos(here)
        return None

    @contextual
    def output(self) -> OutputEffect | None:
        return self.stack_effect()

    @contextual
    def cache_effect(self) -> CacheEffect | None:
        
        if tkn := self.expect(lx.IDENTIFIER):
            if self.expect(lx.DIVIDE):
                num = self.require(lx.NUMBER).text
                try:
                    size = int(num)
                except ValueError:
                    raise self.make_syntax_error(f"Expected integer, got {num!r}")
                else:
                    return CacheEffect(tkn.text, size)
        return None

    @contextual
    def stack_effect(self) -> StackEffect | None:
        
        
        if tkn := self.expect(lx.IDENTIFIER):
            type_text = ""
            if self.expect(lx.COLON):
                type_text = self.require(lx.IDENTIFIER).text.strip()
                if self.expect(lx.TIMES):
                    type_text += " *"
            cond_text = ""
            if self.expect(lx.IF):
                self.require(lx.LPAREN)
                if not (cond := self.expression()):
                    raise self.make_syntax_error("Expected condition")
                self.require(lx.RPAREN)
                cond_text = cond.text.strip()
            size_text = ""
            if self.expect(lx.LBRACKET):
                if type_text or cond_text:
                    raise self.make_syntax_error("Unexpected [")
                if not (size := self.expression()):
                    raise self.make_syntax_error("Expected expression")
                self.require(lx.RBRACKET)
                type_text = "PyObject **"
                size_text = size.text.strip()
            return StackEffect(tkn.text, type_text, cond_text, size_text)
        return None

    @contextual
    def expression(self) -> Expression | None:
        tokens: list[lx.Token] = []
        level = 1
        while tkn := self.peek():
            if tkn.kind in (lx.LBRACKET, lx.LPAREN):
                level += 1
            elif tkn.kind in (lx.RBRACKET, lx.RPAREN):
                level -= 1
                if level == 0:
                    break
            tokens.append(tkn)
            self.next()
        if not tokens:
            return None
        return Expression(lx.to_text(tokens).strip())

    
    
    
    
    
    
    

    @contextual
    def op(self) -> OpName | None:
        if tkn := self.expect(lx.IDENTIFIER):
            return OpName(tkn.text)
        return None

    @contextual
    def macro_def(self) -> Macro | None:
        if tkn := self.expect(lx.MACRO):
            if self.expect(lx.LPAREN):
                if tkn := self.expect(lx.IDENTIFIER):
                    if self.expect(lx.RPAREN):
                        if self.expect(lx.EQUALS):
                            if uops := self.uops():
                                self.require(lx.SEMI)
                                res = Macro(tkn.text, uops)
                                return res
        return None

    def uops(self) -> list[UOp] | None:
        if uop := self.uop():
            uop = cast(UOp, uop)
            uops = [uop]
            while self.expect(lx.PLUS):
                if uop := self.uop():
                    uop = cast(UOp, uop)
                    uops.append(uop)
                else:
                    raise self.make_syntax_error("Expected op name or cache effect")
            return uops
        return None

    @contextual
    def uop(self) -> UOp | None:
        if tkn := self.expect(lx.IDENTIFIER):
            if self.expect(lx.DIVIDE):
                if num := self.expect(lx.NUMBER):
                    try:
                        size = int(num.text)
                    except ValueError:
                        raise self.make_syntax_error(
                            f"Expected integer, got {num.text!r}"
                        )
                    else:
                        return CacheEffect(tkn.text, size)
                raise self.make_syntax_error("Expected integer")
            else:
                return OpName(tkn.text)
        return None

    @contextual
    def family_def(self) -> Family | None:
        if (tkn := self.expect(lx.IDENTIFIER)) and tkn.text == "family":
            size = None
            if self.expect(lx.LPAREN):
                if tkn := self.expect(lx.IDENTIFIER):
                    if self.expect(lx.COMMA):
                        if not (size := self.expect(lx.IDENTIFIER)):
                            if not (size := self.expect(lx.NUMBER)):
                                raise self.make_syntax_error("Expected identifier or number")
                    if self.expect(lx.RPAREN):
                        if self.expect(lx.EQUALS):
                            if not self.expect(lx.LBRACE):
                                raise self.make_syntax_error("Expected {")
                            if members := self.members():
                                if self.expect(lx.RBRACE) and self.expect(lx.SEMI):
                                    return Family(
                                        tkn.text, size.text if size else "", members
                                    )
        return None

    def flags(self) -> list[str]:
        here = self.getpos()
        if self.expect(lx.LPAREN):
            if tkn := self.expect(lx.IDENTIFIER):
                flags = [tkn.text]
                while self.expect(lx.COMMA):
                    if tkn := self.expect(lx.IDENTIFIER):
                        flags.append(tkn.text)
                    else:
                        break
                if not self.expect(lx.RPAREN):
                    raise self.make_syntax_error("Expected comma or right paren")
                return flags
        self.setpos(here)
        return []

    @contextual
    def pseudo_def(self) -> Pseudo | None:
        if (tkn := self.expect(lx.IDENTIFIER)) and tkn.text == "pseudo":
            size = None
            if self.expect(lx.LPAREN):
                if tkn := self.expect(lx.IDENTIFIER):
                    if self.expect(lx.COMMA):
                        flags = self.flags()
                    else:
                        flags = []
                    if self.expect(lx.RPAREN):
                        if self.expect(lx.EQUALS):
                            if not self.expect(lx.LBRACE):
                                raise self.make_syntax_error("Expected {")
                            if members := self.members():
                                if self.expect(lx.RBRACE) and self.expect(lx.SEMI):
                                    return Pseudo(tkn.text, flags, members)
        return None

    def members(self) -> list[str] | None:
        here = self.getpos()
        if tkn := self.expect(lx.IDENTIFIER):
            members = [tkn.text]
            while self.expect(lx.COMMA):
                if tkn := self.expect(lx.IDENTIFIER):
                    members.append(tkn.text)
                else:
                    break
            peek = self.peek()
            if not peek or peek.kind != lx.RBRACE:
                raise self.make_syntax_error("Expected comma or right paren")
            return members
        self.setpos(here)
        return None

    @contextual
    def block(self) -> Block | None:
        if self.c_blob():
            return Block()
        return None

    def c_blob(self) -> list[lx.Token]:
        tokens: list[lx.Token] = []
        level = 0
        while tkn := self.next(raw=True):
            tokens.append(tkn)
            if tkn.kind in (lx.LBRACE, lx.LPAREN, lx.LBRACKET):
                level += 1
            elif tkn.kind in (lx.RBRACE, lx.RPAREN, lx.RBRACKET):
                level -= 1
                if level <= 0:
                    break
        return tokens


if __name__ == "__main__":
    import sys

    if sys.argv[1:]:
        filename = sys.argv[1]
        if filename == "-c" and sys.argv[2:]:
            src = sys.argv[2]
            filename = "<string>"
        else:
            with open(filename, "r") as f:
                src = f.read()
            srclines = src.splitlines()
            begin = srclines.index("// BEGIN BYTECODES //")
            end = srclines.index("// END BYTECODES //")
            src = "\n".join(srclines[begin + 1 : end])
    else:
        filename = "<default>"
        src = "if (x) { x.foo; // comment\n}"
    parser = Parser(src, filename)
    x = parser.definition()
    print(x)