parser.py

"""This is the NILE parser"""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing_extensions import TypeGuard

from typing import List, Literal, Optional, Tuple, NamedTuple

from lark import Lark, Transformer
from lark.visitors import v_args

parser = Lark(
    r"""
    start: statement

    DATE: /\d{4}\-\d{2}\-\d{2}T\d{2}\:\d{2}\:\d{2}(Z|\+\d{2}\:\d{2})/
    IDENT: /[a-zA-Z][a-zA-Z0-9_\-]*/
    STRING: /[^']+/

    _string: "'" STRING "'"

    _qident: "'" IDENT "'"

    statement: "define" "intent" IDENT ":" _command+

    _command: action
            | target
            | endpoints
            | path
            | period
            | middlebox

    action: "do" "action" "(" _qident ")"

    target: "for" "target" "(" _qident ")"

    endpoints: "from" _endpoint "to" _endpoint
    _endpoint: "endpoint" "(" _qident ")"

    path: "following" "path" "(" _qident ("," _qident)* ")"
    period: "start" _time "to" _time
    _time: "hour" "(" "'" DATE "'" ")"

    middlebox: "add" _mboxname ("," _mboxname)* ("with" middlebox_qos)* middlebox_block*
    _mboxname: "middlebox" "(" _qident ")"
    middlebox_qos: /latency|throughput|jitter|loss/ "(" (_string | "'" "none" "'") ["," _string] ")"
    middlebox_block: "block" "traffic" "(" _qident ")"

    %import common.WS
    %ignore WS
""",
    parser="lalr",
)


class NileAction(NamedTuple):
    name: str


class NileTarget(NamedTuple):
    name: str


class NileEndpoints(NamedTuple):
    src: str
    dst: str


NileQOSType = Literal["latency", "throughput", "jitter", "loss"]


class NileQOS(NamedTuple):
    type: NileQOSType
    constraint: Optional[Tuple[str, str]]


class NileBlock(NamedTuple):
    name: str


class NileMiddlebox(NamedTuple):
    name: str
    qos: List[NileQOS]
    blocks: List[NileBlock]


class NilePath(NamedTuple):
    path: List[str]


class NilePeriod(NamedTuple):
    start: str  # TODO: parse
    end: str


# TODO: should we really just throw everything into a list


class NileIntent(NamedTuple):
    name: str
    actions: List[NileAction]
    targets: List[NileTarget]
    endpoints: List[NileEndpoints]
    paths: List[NilePath]
    periods: List[NilePeriod]
    middleboxes: List[NileMiddlebox]


def ensure_middlebox_type(s: str) -> TypeGuard[NileQOSType]:
    return s in ("latency", "throughput", "jitter", "loss")


class NileTransformer(Transformer):
    def start(self, args):
        return args[0]

    def statement(self, args):
        name = args[0]

        stmt = NileIntent(name, [], [], [], [], [], [])

        attrs = {
            NileAction: stmt.actions,
            NileTarget: stmt.targets,
            NileEndpoints: stmt.endpoints,
            NilePath: stmt.paths,
            NilePeriod: stmt.periods,
            NileMiddlebox: stmt.middleboxes,
        }

        for command in args[1:]:
            attrs[type(command)].append(command)

        return stmt

    DATE = str
    IDENT = str
    STRING = str

    action = v_args(inline=True)(NileAction)
    target = v_args(inline=True)(NileTarget)
    endpoints = v_args(inline=True)(NileEndpoints)
    path = NilePath
    period = v_args(inline=True)(NilePeriod)

    def middlebox_qos(self, args):
        type_, *params = args
        type_ = str(type_)
        assert ensure_middlebox_type(type_)

        constraint = (params[0], params[1]) if params else None
        return NileQOS(type_, constraint)

    middlebox_block = v_args(inline=True)(NileBlock)

    def middlebox(self, args):
        name, *rest = args
        mbox = NileMiddlebox(name, [], [])

        attrs = {
            NileQOS: mbox.qos,
            NileBlock: mbox.blocks,
        }

        for thing in rest:
            attrs[type(thing)].append(thing)

        return mbox


def parse(input: str) -> NileIntent:
    r = parser.parse(input)
    return NileTransformer().transform(r)