lib9p: protogen: c_marshal.py: Clean up

1 files changed, 261 insertions, 241 deletions

diff --git a/lib9p/protogen/c_marshal.py b/lib9p/protogen/c_marshal.py
index 152206d..38a2feb 100644
--- a/lib9p/protogen/c_marshal.py
+++ b/lib9p/protogen/c_marshal.py
@@ -17,6 +17,162 @@ from . import c9util, cutil, idlutil
 # pylint: disable=unused-variable
 __all__ = ["gen_c_marshal"]
 
+# get_offset_expr() ############################################################
+
+
+class OffsetExpr:
+    static: int
+    cond: dict[frozenset[str], "OffsetExpr"]
+    rep: list[tuple[idlutil.Path, "OffsetExpr"]]
+
+    def __init__(self) -> None:
+        self.static = 0
+        self.rep = []
+        self.cond = {}
+
+    def add(self, other: "OffsetExpr") -> None:
+        self.static += other.static
+        self.rep += other.rep
+        for k, v in other.cond.items():
+            if k in self.cond:
+                self.cond[k].add(v)
+            else:
+                self.cond[k] = v
+
+    def gen_c(
+        self,
+        dsttyp: str,
+        dstvar: str,
+        root: str,
+        indent_depth: int,
+        loop_depth: int,
+    ) -> str:
+        oneline: list[str] = []
+        multiline = ""
+        if self.static:
+            oneline.append(str(self.static))
+        for cnt, sub in self.rep:
+            if not sub.cond and not sub.rep:
+                if sub.static == 1:
+                    oneline.append(cnt.c_str(root))
+                else:
+                    oneline.append(f"({cnt.c_str(root)})*{sub.static}")
+                continue
+            loopvar = chr(ord("i") + loop_depth)
+            multiline += f"{'\t'*indent_depth}for ({c9util.typename(cnt.elems[-1].typ)} {loopvar} = 0; {loopvar} < {cnt.c_str(root)}; {loopvar}++) {{\n"
+            multiline += sub.gen_c("", dstvar, root, indent_depth + 1, loop_depth + 1)
+            multiline += f"{'\t'*indent_depth}}}\n"
+        for vers, sub in self.cond.items():
+            multiline += cutil.ifdef_push(indent_depth + 1, c9util.ver_ifdef(vers))
+            multiline += f"{'\t'*indent_depth}if {c9util.ver_cond(vers)} {{\n"
+            multiline += sub.gen_c("", dstvar, root, indent_depth + 1, loop_depth)
+            multiline += f"{'\t'*indent_depth}}}\n"
+            multiline += cutil.ifdef_pop(indent_depth)
+        ret = ""
+        if dsttyp:
+            if not oneline:
+                oneline.append("0")
+            ret += f"{'\t'*indent_depth}{dsttyp} {dstvar} = {' + '.join(oneline)};\n"
+        elif oneline:
+            ret += f"{'\t'*indent_depth}{dstvar} += {' + '.join(oneline)};\n"
+        ret += multiline
+        return ret
+
+
+type OffsetExprRecursion = typing.Callable[[idlutil.Path], idlutil.WalkCmd]
+
+
+def get_offset_expr(typ: idl.UserType, recurse: OffsetExprRecursion) -> OffsetExpr:
+    if not isinstance(typ, idl.Struct):
+        assert typ.static_size
+        ret = OffsetExpr()
+        ret.static = typ.static_size
+        return ret
+
+    class ExprStackItem(typing.NamedTuple):
+        path: idlutil.Path
+        expr: OffsetExpr
+        pop: typing.Callable[[], None]
+
+    expr_stack: list[ExprStackItem]
+
+    def pop_root() -> None:
+        assert False
+
+    def pop_cond() -> None:
+        nonlocal expr_stack
+        key = frozenset(expr_stack[-1].path.elems[-1].in_versions)
+        if key in expr_stack[-2].expr.cond:
+            expr_stack[-2].expr.cond[key].add(expr_stack[-1].expr)
+        else:
+            expr_stack[-2].expr.cond[key] = expr_stack[-1].expr
+        expr_stack = expr_stack[:-1]
+
+    def pop_rep() -> None:
+        nonlocal expr_stack
+        member_path = expr_stack[-1].path
+        member = member_path.elems[-1]
+        assert member.cnt
+        cnt_path = member_path.parent().add(member.cnt)
+        expr_stack[-2].expr.rep.append((cnt_path, expr_stack[-1].expr))
+        expr_stack = expr_stack[:-1]
+
+    def handle(
+        path: idlutil.Path,
+    ) -> tuple[idlutil.WalkCmd, typing.Callable[[], None] | None]:
+        nonlocal recurse
+
+        ret = recurse(path)
+        if ret != idlutil.WalkCmd.KEEP_GOING:
+            return ret, None
+
+        nonlocal expr_stack
+        expr_stack_len = len(expr_stack)
+
+        def pop() -> None:
+            nonlocal expr_stack
+            nonlocal expr_stack_len
+            while len(expr_stack) > expr_stack_len:
+                expr_stack[-1].pop()
+
+        if path.elems:
+            child = path.elems[-1]
+            parent = path.elems[-2].typ if len(path.elems) > 1 else path.root
+            if child.in_versions < parent.in_versions:
+                expr_stack.append(
+                    ExprStackItem(path=path, expr=OffsetExpr(), pop=pop_cond)
+                )
+            if child.cnt:
+                expr_stack.append(
+                    ExprStackItem(path=path, expr=OffsetExpr(), pop=pop_rep)
+                )
+            if not isinstance(child.typ, idl.Struct):
+                assert child.typ.static_size
+                expr_stack[-1].expr.static += child.typ.static_size
+        return ret, pop
+
+    expr_stack = [
+        ExprStackItem(path=idlutil.Path(typ), expr=OffsetExpr(), pop=pop_root)
+    ]
+    idlutil.walk(typ, handle)
+    return expr_stack[0].expr
+
+
+def go_to_end(path: idlutil.Path) -> idlutil.WalkCmd:
+    return idlutil.WalkCmd.KEEP_GOING
+
+
+def go_to_tok(name: str) -> typing.Callable[[idlutil.Path], idlutil.WalkCmd]:
+    def ret(path: idlutil.Path) -> idlutil.WalkCmd:
+        if len(path.elems) == 1 and path.elems[0].membname == name:
+            return idlutil.WalkCmd.ABORT
+        return idlutil.WalkCmd.KEEP_GOING
+
+    return ret
+
+
+# Generate .c ##################################################################
+
 
 def gen_c_marshal(versions: set[str], typs: list[idl.UserType]) -> str:
     ret = """
@@ -72,141 +228,112 @@ def gen_c_marshal(versions: set[str], typs: list[idl.UserType]) -> str:
         "\tctx->net_iov[ctx->net_iov_cnt-1].iov_len += 8;\n"
     )
 
-    class OffsetExpr:
-        static: int
-        cond: dict[frozenset[str], "OffsetExpr"]
-        rep: list[tuple[idlutil.Path, "OffsetExpr"]]
-
-        def __init__(self) -> None:
-            self.static = 0
-            self.rep = []
-            self.cond = {}
-
-        def add(self, other: "OffsetExpr") -> None:
-            self.static += other.static
-            self.rep += other.rep
-            for k, v in other.cond.items():
-                if k in self.cond:
-                    self.cond[k].add(v)
-                else:
-                    self.cond[k] = v
-
-        def gen_c(
-            self,
-            dsttyp: str,
-            dstvar: str,
-            root: str,
-            indent_depth: int,
-            loop_depth: int,
-        ) -> str:
-            oneline: list[str] = []
-            multiline = ""
-            if self.static:
-                oneline.append(str(self.static))
-            for cnt, sub in self.rep:
-                if not sub.cond and not sub.rep:
-                    if sub.static == 1:
-                        oneline.append(cnt.c_str(root))
-                    else:
-                        oneline.append(f"({cnt.c_str(root)})*{sub.static}")
-                    continue
-                loopvar = chr(ord("i") + loop_depth)
-                multiline += f"{'\t'*indent_depth}for ({c9util.typename(cnt.elems[-1].typ)} {loopvar} = 0; {loopvar} < {cnt.c_str(root)}; {loopvar}++) {{\n"
-                multiline += sub.gen_c(
-                    "", dstvar, root, indent_depth + 1, loop_depth + 1
-                )
-                multiline += f"{'\t'*indent_depth}}}\n"
-            for vers, sub in self.cond.items():
-                multiline += cutil.ifdef_push(indent_depth + 1, c9util.ver_ifdef(vers))
-                multiline += f"{'\t'*indent_depth}if {c9util.ver_cond(vers)} {{\n"
-                multiline += sub.gen_c("", dstvar, root, indent_depth + 1, loop_depth)
-                multiline += f"{'\t'*indent_depth}}}\n"
-                multiline += cutil.ifdef_pop(indent_depth)
-            if dsttyp:
-                if not oneline:
-                    oneline.append("0")
-                ret = f"{'\t'*indent_depth}{dsttyp} {dstvar} = {' + '.join(oneline)};\n"
-            elif oneline:
-                ret = f"{'\t'*indent_depth}{dstvar} += {' + '.join(oneline)};\n"
-            ret += multiline
-            return ret
-
-    type OffsetExprRecursion = typing.Callable[[idlutil.Path], idlutil.WalkCmd]
-
-    def get_offset_expr(typ: idl.UserType, recurse: OffsetExprRecursion) -> OffsetExpr:
-        if not isinstance(typ, idl.Struct):
-            assert typ.static_size
-            ret = OffsetExpr()
-            ret.static = typ.static_size
-            return ret
-
-        stack: list[tuple[idlutil.Path, OffsetExpr, typing.Callable[[], None]]]
-
-        def pop_root() -> None:
-            assert False
-
-        def pop_cond() -> None:
-            nonlocal stack
-            key = frozenset(stack[-1][0].elems[-1].in_versions)
-            if key in stack[-2][1].cond:
-                stack[-2][1].cond[key].add(stack[-1][1])
-            else:
-                stack[-2][1].cond[key] = stack[-1][1]
-            stack = stack[:-1]
-
-        def pop_rep() -> None:
-            nonlocal stack
-            member_path = stack[-1][0]
-            member = member_path.elems[-1]
-            assert member.cnt
-            cnt_path = member_path.parent().add(member.cnt)
-            stack[-2][1].rep.append((cnt_path, stack[-1][1]))
-            stack = stack[:-1]
-
-        def handle(
-            path: idlutil.Path,
-        ) -> tuple[idlutil.WalkCmd, typing.Callable[[], None] | None]:
-            nonlocal recurse
-
-            ret = recurse(path)
-            if ret != idlutil.WalkCmd.KEEP_GOING:
-                return ret, None
-
-            nonlocal stack
-            stack_len = len(stack)
-
-            def pop() -> None:
-                nonlocal stack
-                nonlocal stack_len
-                while len(stack) > stack_len:
-                    stack[-1][2]()
-
-            if path.elems:
-                child = path.elems[-1]
-                parent = path.elems[-2].typ if len(path.elems) > 1 else path.root
-                if child.in_versions < parent.in_versions:
-                    stack.append((path, OffsetExpr(), pop_cond))
-                if child.cnt:
-                    stack.append((path, OffsetExpr(), pop_rep))
-                if not isinstance(child.typ, idl.Struct):
-                    assert child.typ.static_size
-                    stack[-1][1].static += child.typ.static_size
-            return ret, pop
-
-        stack = [(idlutil.Path(typ), OffsetExpr(), pop_root)]
-        idlutil.walk(typ, handle)
-        return stack[0][1]
+    class IndentLevel(typing.NamedTuple):
+        ifdef: bool  # whether this is both `{` and `#if`, or just `{`
 
-    def go_to_end(path: idlutil.Path) -> idlutil.WalkCmd:
-        return idlutil.WalkCmd.KEEP_GOING
+    indent_stack: list[IndentLevel]
 
-    def go_to_tok(name: str) -> typing.Callable[[idlutil.Path], idlutil.WalkCmd]:
-        def ret(path: idlutil.Path) -> idlutil.WalkCmd:
-            if len(path.elems) == 1 and path.elems[0].membname == name:
-                return idlutil.WalkCmd.ABORT
-            return idlutil.WalkCmd.KEEP_GOING
+    def ifdef_lvl() -> int:
+        return sum(1 if lvl.ifdef else 0 for lvl in indent_stack)
 
-        return ret
+    def indent_lvl() -> int:
+        return len(indent_stack)
+
+    max_size: int
+
+    def handle(
+        path: idlutil.Path,
+    ) -> tuple[idlutil.WalkCmd, typing.Callable[[], None]]:
+        nonlocal ret
+        nonlocal indent_stack
+        nonlocal max_size
+        indent_stack_len = len(indent_stack)
+
+        def pop() -> None:
+            nonlocal ret
+            nonlocal indent_stack
+            nonlocal indent_stack_len
+            while len(indent_stack) > indent_stack_len:
+                ret += f"{'\t'*(indent_lvl()-1)}}}\n"
+                if indent_stack.pop().ifdef:
+                    ret += cutil.ifdef_pop(ifdef_lvl())
+
+        loopdepth = sum(1 for elem in path.elems if elem.cnt)
+        struct = path.elems[-1].typ if path.elems else path.root
+        if isinstance(struct, idl.Struct):
+            offsets: list[str] = []
+            for member in struct.members:
+                if not member.val:
+                    continue
+                for tok in member.val.tokens:
+                    if not isinstance(tok, idl.ExprSym):
+                        continue
+                    if tok.symname == "end" or tok.symname.startswith("&"):
+                        if tok.symname not in offsets:
+                            offsets.append(tok.symname)
+            for name in offsets:
+                name_prefix = f"offsetof{''.join('_'+m.membname for m in path.elems)}_"
+                if name == "end":
+                    if not path.elems:
+                        if max_size > cutil.UINT32_MAX:
+                            ret += f"{'\t'*indent_lvl()}uint32_t {name_prefix}end = (uint32_t)needed_size;\n"
+                        else:
+                            ret += f"{'\t'*indent_lvl()}uint32_t {name_prefix}end = needed_size;\n"
+                        continue
+                    recurse: OffsetExprRecursion = go_to_end
+                else:
+                    assert name.startswith("&")
+                    name = name[1:]
+                    recurse = go_to_tok(name)
+                expr = get_offset_expr(struct, recurse)
+                expr_prefix = path.c_str("val->", loopdepth)
+                if not expr_prefix.endswith(">"):
+                    expr_prefix += "."
+                ret += expr.gen_c(
+                    "uint32_t",
+                    name_prefix + name,
+                    expr_prefix,
+                    indent_lvl(),
+                    loopdepth,
+                )
+        if not path.elems:
+            return idlutil.WalkCmd.KEEP_GOING, pop
+
+        child = path.elems[-1]
+        parent = path.elems[-2].typ if len(path.elems) > 1 else path.root
+        if child.in_versions < parent.in_versions:
+            ret += cutil.ifdef_push(
+                ifdef_lvl() + 1, c9util.ver_ifdef(child.in_versions)
+            )
+            ret += f"{'\t'*indent_lvl()}if ({c9util.ver_cond(child.in_versions)}) {{\n"
+            indent_stack.append(IndentLevel(ifdef=True))
+        if child.cnt:
+            cnt_path = path.parent().add(child.cnt)
+            if child.typ.static_size == 1:  # SPECIAL (zerocopy)
+                if path.root.typname == "stat":  # SPECIAL (stat)
+                    ret += f"{'\t'*indent_lvl()}MARSHAL_BYTES(ctx, {path.c_str('val->')[:-3]}, {cnt_path.c_str('val->')});\n"
+                else:
+                    ret += f"{'\t'*indent_lvl()}MARSHAL_BYTES_ZEROCOPY(ctx, {path.c_str('val->')[:-3]}, {cnt_path.c_str('val->')});\n"
+                return idlutil.WalkCmd.KEEP_GOING, pop
+            loopvar = chr(ord("i") + loopdepth - 1)
+            ret += f"{'\t'*indent_lvl()}for ({c9util.typename(child.cnt.typ)} {loopvar} = 0; {loopvar} < {cnt_path.c_str('val->')}; {loopvar}++) {{\n"
+            indent_stack.append(IndentLevel(ifdef=False))
+        if not isinstance(child.typ, idl.Struct):
+            if child.val:
+
+                def lookup_sym(sym: str) -> str:
+                    nonlocal path
+                    if sym.startswith("&"):
+                        sym = sym[1:]
+                    return f"offsetof{''.join('_'+m.membname for m in path.elems[:-1])}_{sym}"
+
+                val = c9util.idl_expr(child.val, lookup_sym)
+            else:
+                val = path.c_str("val->")
+            if isinstance(child.typ, idl.Bitfield):
+                val += f" & {child.typ.typname}_masks[ctx->ctx->version]"
+            ret += f"{'\t'*indent_lvl()}MARSHAL_U{child.typ.static_size*8}LE(ctx, {val});\n"
+        return idlutil.WalkCmd.KEEP_GOING, pop
 
     for typ in typs:
         if not (
@@ -240,117 +367,10 @@ def gen_c_marshal(versions: set[str], typs: list[idl.UserType]) -> str:
         ret += "\t}\n"
 
         # Pass 2 - write data
-        ifdef_depth = 1
-        stack: list[tuple[idlutil.Path, bool]] = [(idlutil.Path(typ), False)]
-
-        def handle(
-            path: idlutil.Path,
-        ) -> tuple[idlutil.WalkCmd, typing.Callable[[], None]]:
-            nonlocal ret
-            nonlocal ifdef_depth
-            nonlocal stack
-            stack_len = len(stack)
-
-            def pop() -> None:
-                nonlocal ret
-                nonlocal ifdef_depth
-                nonlocal stack
-                nonlocal stack_len
-                while len(stack) > stack_len:
-                    ret += f"{'\t'*(len(stack)-1)}}}\n"
-                    if stack[-1][1]:
-                        ifdef_depth -= 1
-                        ret += cutil.ifdef_pop(ifdef_depth)
-                    stack = stack[:-1]
-
-            loopdepth = sum(1 for elem in path.elems if elem.cnt)
-            struct = path.elems[-1].typ if path.elems else path.root
-            if isinstance(struct, idl.Struct):
-                offsets: list[str] = []
-                for member in struct.members:
-                    if not member.val:
-                        continue
-                    for tok in member.val.tokens:
-                        if not isinstance(tok, idl.ExprSym):
-                            continue
-                        if tok.symname == "end" or tok.symname.startswith("&"):
-                            if tok.symname not in offsets:
-                                offsets.append(tok.symname)
-                for name in offsets:
-                    name_prefix = "offsetof_" + "".join(
-                        m.membname + "_" for m in path.elems
-                    )
-                    if name == "end":
-                        if not path.elems:
-                            nonlocal max_size
-                            if max_size > cutil.UINT32_MAX:
-                                ret += f"{'\t'*len(stack)}uint32_t {name_prefix}end = (uint32_t)needed_size;\n"
-                            else:
-                                ret += f"{'\t'*len(stack)}uint32_t {name_prefix}end = needed_size;\n"
-                            continue
-                        recurse: OffsetExprRecursion = go_to_end
-                    else:
-                        assert name.startswith("&")
-                        name = name[1:]
-                        recurse = go_to_tok(name)
-                    expr = get_offset_expr(struct, recurse)
-                    expr_prefix = path.c_str("val->", loopdepth)
-                    if not expr_prefix.endswith(">"):
-                        expr_prefix += "."
-                    ret += expr.gen_c(
-                        "uint32_t",
-                        name_prefix + name,
-                        expr_prefix,
-                        len(stack),
-                        loopdepth,
-                    )
-            if path.elems:
-                child = path.elems[-1]
-                parent = path.elems[-2].typ if len(path.elems) > 1 else path.root
-                if child.in_versions < parent.in_versions:
-                    ret += cutil.ifdef_push(
-                        ifdef_depth + 1, c9util.ver_ifdef(child.in_versions)
-                    )
-                    ifdef_depth += 1
-                    ret += f"{'\t'*len(stack)}if ({c9util.ver_cond(child.in_versions)}) {{\n"
-                    stack.append((path, True))
-                if child.cnt:
-                    cnt_path = path.parent().add(child.cnt)
-                    if child.typ.static_size == 1:  # SPECIAL (zerocopy)
-                        if path.root.typname == "stat":  # SPECIAL (stat)
-                            ret += f"{'\t'*len(stack)}MARSHAL_BYTES(ctx, {path.c_str('val->')[:-3]}, {cnt_path.c_str('val->')});\n"
-                        else:
-                            ret += f"{'\t'*len(stack)}MARSHAL_BYTES_ZEROCOPY(ctx, {path.c_str('val->')[:-3]}, {cnt_path.c_str('val->')});\n"
-                        return idlutil.WalkCmd.KEEP_GOING, pop
-                    loopvar = chr(ord("i") + loopdepth - 1)
-                    ret += f"{'\t'*len(stack)}for ({c9util.typename(child.cnt.typ)} {loopvar} = 0; {loopvar} < {cnt_path.c_str('val->')}; {loopvar}++) {{\n"
-                    stack.append((path, False))
-                if not isinstance(child.typ, idl.Struct):
-                    if child.val:
-
-                        def lookup_sym(sym: str) -> str:
-                            nonlocal path
-                            if sym.startswith("&"):
-                                sym = sym[1:]
-                            return (
-                                "offsetof_"
-                                + "".join(m.membname + "_" for m in path.elems[:-1])
-                                + sym
-                            )
-
-                        val = c9util.idl_expr(child.val, lookup_sym)
-                    else:
-                        val = path.c_str("val->")
-                    if isinstance(child.typ, idl.Bitfield):
-                        val += f" & {child.typ.typname}_masks[ctx->ctx->version]"
-                    ret += f"{'\t'*len(stack)}MARSHAL_U{child.typ.static_size*8}LE(ctx, {val});\n"
-            return idlutil.WalkCmd.KEEP_GOING, pop
-
+        indent_stack = [IndentLevel(ifdef=True)]
         idlutil.walk(typ, handle)
-        del handle
-        del stack
-        del max_size
 
+        # Return
         ret += "\treturn false;\n"
         ret += "}\n"
     ret += cutil.ifdef_pop(0)