loc_notation.py

#!/usr/bin/env python3

from abc import ABC
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from dataclasses import dataclass
from functools import partialmethod
from sys import argv, exit, stderr, stdin
from textwrap import dedent

from lark import Discard, Lark, Transformer, v_args

grammar = """
start: stmt+

?stmt: value
     | var_def

?value: locdesc_value
      | storage_value

?var_def: locdesc_var_def
        | storage_var_def

?locdesc_value: locdesc
              | locdesc_var_ref
locdesc: "(" offset storage_value ")"
locdesc_var_def: "set" "$(" CNAME ")" "=" locdesc
locdesc_var_ref: "$(" CNAME ")"

?storage_value: storage
              | storage_var_ref
storage: "<" storage_body ">"
storage_body: "Undefined" -> undefined
             | "Memory" ID -> memory
             | "Register" ID -> register
             | "ImplicitBytes" INT* -> implicit_bytes
             | "ImplicitValue" ESCAPED_STRING INT -> implicit_value
             | "Composite" (size locdesc)+ -> composite
storage_var_def: "set" "$<" CNAME ">" "=" storage
storage_var_ref: "$<" CNAME ">"

ID: INT | ESCAPED_STRING
size: quantity
offset: quantity
?quantity: BYTE_QUANTITY
         | BIT_QUANTITY
         | BARE_QUANTITY
BYTE_QUANTITY: /[0-9]+B/
BIT_QUANTITY: /[0-9]+b/
BARE_QUANTITY: /[0-9]+/

COMMENT: /;[^\n]*/x

%import common.ESCAPED_STRING
%import common.CNAME
%import common.INT
%import common.WS
%ignore WS
%ignore COMMENT

"""
lark = Lark(grammar)


@v_args(inline=True)
class SubstituteVars(Transformer):
    def __init__(self):
        self.vars = {"locdesc": {}, "storage": {}}

    def _var_def(self, kind, name, value):
        self.vars[kind][name] = value
        return Discard

    def _var_ref(self, kind, name):
        return self.vars[kind][name]

    locdesc_var_def = partialmethod(_var_def, "locdesc")
    locdesc_var_ref = partialmethod(_var_ref, "locdesc")
    storage_var_def = partialmethod(_var_def, "storage")
    storage_var_ref = partialmethod(_var_ref, "storage")


@dataclass
class Statements:
    statements: list

    def __repr__(self):
        return "\n".join(map(repr, self.statements))


class Storage(ABC):
    kind_subclasses = {}

    def __init_subclass__(cls, /, rule, **kwargs):
        super().__init_subclass__(**kwargs)
        if rule in Storage.kind_subclasses:
            raise ValueError(f"redefined storage kind subclass: {rule}")
        Storage.kind_subclasses[rule] = cls

    @classmethod
    def from_children(cls, children):
        assert len(children) == 1
        tree = children[0]
        return cls.kind_subclasses[str(tree.data)](*tree.children)


@dataclass
class LocDesc:
    offset: int
    storage: Storage

    def __repr__(self):
        return f"({self.offset}b {self.storage!r})"


@dataclass
class Undefined(Storage, rule="undefined"):
    def __repr__(self):
        return "<Undefined>"


@dataclass
class Memory(Storage, rule="memory"):
    id: str | int

    def __repr__(self):
        return f"<Memory {self.id}>"


@dataclass
class Register(Storage, rule="register"):
    id: str | int

    def __repr__(self):
        return f"<Register {self.id}>"


class ImplicitBytes(Storage, rule="implicit_bytes"):
    __match_args__ = ("bites",)

    def __init__(self, *args):
        self.bites = list(map(int, args))

    def __repr__(self):
        return f'<ImplicitBytes {" ".join(map(repr, self.bites))}>'


@dataclass
class ImplicitValue(Storage, rule="implicit_value"):
    typeid: str
    value: int

    def __repr__(self):
        return f"<ImplicitValue {self.typeid} {self.value}>"


@dataclass
class Part:
    size: int
    locdesc: LocDesc

    def __repr__(self):
        return f"({self.size}b {self.locdesc!r})"


def _pairsof(iterable):
    it = iter(iterable)
    return zip(it, it, strict=True)


class Composite(Storage, rule="composite"):
    __match_args__ = ("parts",)

    def __init__(self, *args):
        self.parts = [Part(*pair) for pair in _pairsof(args)]

    def __repr__(self):
        return f'<Composite {" ".join(map(repr, self.parts))}>'


class Typeify(Transformer):
    def start(self, c):
        return Statements(c)

    def locdesc(self, c):
        return LocDesc(c[0], c[1])

    def storage(self, c):
        return Storage.from_children(c)

    def offset(self, c):
        return c[0]

    def size(self, c):
        return c[0]

    def BYTE_QUANTITY(self, c):
        return int(c[:-1]) * 8

    def BIT_QUANTITY(self, c):
        return int(c[:-1])

    def BARE_QUANTITY(self, c):
        return int(c)


def parse(s):
    return Typeify().transform(SubstituteVars().transform(lark.parse(s)))


def none(v):
    return repr(v)


def email(v):
    match v:
        case Statements(statements):
            return "\n\n".join(map(email, statements))
        case LocDesc(offset, Composite(parts)):
            s = f'|-- {" --|-- ".join(map(lambda part: email(part.locdesc), parts))} --|'
            p = ""
            fill = lambda: s.find("|", len(p)) - len(p)
            part_offset = 0
            for part in parts:
                p += "." * fill() + str(part_offset)
                part_offset += part.size
            p += "." * (fill() - len(str(part_offset)) + 1) + str(part_offset)
            return f"  parts: {p}\nstorage: {s}\n offset: {offset}"
        case LocDesc(offset, storage):
            return f"{email(storage)} (offset: {offset})"
        case Undefined():
            return f"Undef"
        case Memory(id):
            return f"Mem({id})"
        case Register(id):
            return f"R{id}"
        case ImplicitBytes(bites):
            return f"Implicit{bites}"
        case ImplicitValue(typeid, value):
            return f"Implicit({typeid} {value})"
        case _:
            raise ValueError(f"unhandled case {v}")


def verbose(v, depth=1):
    match v:
        case Statements(statements):
            return "\n\n".join(map(lambda s: verbose(s, depth), statements))
        case LocDesc(offset, Composite(parts)):
            indent = "\t" * depth
            parts = ", ".join(
                map(
                    lambda p: f"\n{indent}Part{{size:{p.size}b locdesc:{verbose(p.locdesc, depth + 1)}}}",
                    parts,
                )
            )
            return f"(offset:{offset}b storage:Composite{{parts:[{parts}\n{indent}]}})"
        case LocDesc(offset, storage):
            return f"(offset:{offset}b storage:{verbose(storage, depth)})"
        case Undefined():
            return "Undefined"
        case Memory(id):
            return f"Memory{{addrspace:{id}}}"
        case Register(id):
            return f"Register{{number:{id}}}"
        case ImplicitBytes(bites):
            return f"Implicit{{bytes:{bites}}}"
        case ImplicitValue(typeid, value):
            return f"Implicit{{type:{typeid} value:{value}}}"
        case _:
            return ValueError(f"unhandled case {v}")


renderers = {"none": none, "email": email, "verbose": verbose}

if __name__ == "__main__":
    arg_parser = ArgumentParser(
        formatter_class=RawDescriptionHelpFormatter,
        epilog=dedent(
            f"""
            examples:
              printf "(1 <Undefined>)" | {argv[0]} email
              {argv[0]} email <file
            """
        ),
    )
    arg_parser.add_argument("renderer", help="renderer", choices=renderers.keys())
    args = arg_parser.parse_args()
    print(renderers[args.renderer](parse(stdin.read())))