lib9p: idl: Allow for const .cnt

5 files changed, 65 insertions, 28 deletions

diff --git a/lib9p/idl/0000-TODO.md b/lib9p/idl/0000-TODO.md
index 81cafe5..5079a1d 100644
--- a/lib9p/idl/0000-TODO.md
+++ b/lib9p/idl/0000-TODO.md
@@ -10,4 +10,3 @@
 - Decide how to handle duplicate `enum lib9p_msg_type` names and
   values
 - Clean up the iterate-over-all-msgids-in-a-version code
-- Allow for const `.cnt` instead of only having previous-member `.cnt`
diff --git a/lib9p/idl/__init__.py b/lib9p/idl/__init__.py
index e7b3670..f8e1a38 100644
--- a/lib9p/idl/__init__.py
+++ b/lib9p/idl/__init__.py
@@ -163,7 +163,7 @@ class Expr:
 
 class StructMember:
     # from left-to-right when parsing
-    cnt: "StructMember | None" = None
+    cnt: "StructMember| int | None" = None
     membname: str
     typ: "Type"
     max: Expr
@@ -174,6 +174,8 @@ class StructMember:
     @property
     def min_cnt(self) -> int:
         assert self.cnt
+        if isinstance(self.cnt, int):
+            return self.cnt
         if not isinstance(self.cnt.typ, Primitive):
             raise ValueError(
                 f"list count must be an integer type: {self.cnt.membname!r}"
@@ -185,6 +187,8 @@ class StructMember:
     @property
     def max_cnt(self) -> int:
         assert self.cnt
+        if isinstance(self.cnt, int):
+            return self.cnt
         if not isinstance(self.cnt.typ, Primitive):
             raise ValueError(
                 f"list count must be an integer type: {self.cnt.membname!r}"
@@ -305,7 +309,7 @@ re_bitspec_bit = (
 )
 re_bitspec_alias = f"(?P<name>{re_symname_u})\\s*=\\s*(?P<val>\\S+)"
 
-re_memberspec = f"(?:(?P<cnt>{re_symname})\\*\\()?(?P<name>{re_symname})\\[(?P<typ>{re_memtype})(?:,max=(?P<max>{re_expr})|,val=(?P<val>{re_expr}))*\\]\\)?"
+re_memberspec = f"(?:(?P<cnt>{re_symname}|[1-9][0-9]*)\\*\\()?(?P<name>{re_symname})\\[(?P<typ>{re_memtype})(?:,max=(?P<max>{re_expr})|,val=(?P<val>{re_expr}))*\\]\\)?"
 
 
 def parse_numspec(ver: str, n: Number, spec: str) -> None:
@@ -401,11 +405,13 @@ def parse_members(ver: str, env: dict[str, Type], struct: Struct, specs: str) ->
         member.typ = env[m.group("typ")]
 
         if cnt := m.group("cnt"):
-            if len(struct.members) == 0 or struct.members[-1].membname != cnt:
-                raise ValueError(f"list count must be previous item: {cnt!r}")
-            cnt_mem = struct.members[-1]
-            member.cnt = cnt_mem
-            _ = member.max_cnt  # force validation
+            if cnt.isnumeric():
+                member.cnt = int(cnt)
+            else:
+                if len(struct.members) == 0 or struct.members[-1].membname != cnt:
+                    raise ValueError(f"list count must be previous item: {cnt!r}")
+                member.cnt = struct.members[-1]
+                _ = member.max_cnt  # force validation
 
         if maxstr := m.group("max"):
             if (not isinstance(member.typ, Primitive)) or member.cnt:
diff --git a/lib9p/protogen/c_marshal.py b/lib9p/protogen/c_marshal.py
index 74b64f5..a358cf2 100644
--- a/lib9p/protogen/c_marshal.py
+++ b/lib9p/protogen/c_marshal.py
@@ -23,7 +23,7 @@ __all__ = ["gen_c_marshal"]
 class OffsetExpr:
     static: int
     cond: dict[frozenset[str], "OffsetExpr"]
-    rep: list[tuple[idlutil.Path, "OffsetExpr"]]
+    rep: list[tuple[idlutil.Path | int, "OffsetExpr"]]
 
     def __init__(self) -> None:
         self.static = 0
@@ -52,14 +52,20 @@ class OffsetExpr:
         if self.static:
             oneline.append(str(self.static))
         for cnt, sub in self.rep:
+            if isinstance(cnt, int):
+                cnt_str = str(cnt)
+                cnt_typ = "size_t"
+            else:
+                cnt_str = cnt.c_str(root)
+                cnt_typ = c9util.typename(cnt.elems[-1].typ)
             if not sub.cond and not sub.rep:
                 if sub.static == 1:
-                    oneline.append(cnt.c_str(root))
+                    oneline.append(cnt_str)
                 else:
-                    oneline.append(f"({cnt.c_str(root)})*{sub.static}")
+                    oneline.append(f"({cnt_str})*{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 += f"{'\t'*indent_depth}for ({cnt_typ} {loopvar} = 0; {loopvar} < {cnt_str}; {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():
@@ -113,8 +119,12 @@ def get_offset_expr(typ: idl.UserType, recurse: OffsetExprRecursion) -> OffsetEx
         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))
+        cnt: idlutil.Path | int
+        if isinstance(member.cnt, int):
+            cnt = member.cnt
+        else:
+            cnt = member_path.parent().add(member.cnt)
+        expr_stack[-2].expr.rep.append((cnt, expr_stack[-1].expr))
         expr_stack = expr_stack[:-1]
 
     def handle(
@@ -313,15 +323,20 @@ def gen_c_marshal(versions: set[str], typs: list[idl.UserType]) -> str:
                 )
                 indent_stack.append(IndentLevel(ifdef=True))
         if child.cnt:
-            cnt_path = path.parent().add(child.cnt)
+            if isinstance(child.cnt, int):
+                cnt_str = str(child.cnt)
+                cnt_typ = "size_t"
+            else:
+                cnt_str = path.parent().add(child.cnt).c_str("val->")
+                cnt_typ = c9util.typename(child.cnt.typ)
             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"
+                    ret += f"{'\t'*indent_lvl()}MARSHAL_BYTES(ctx, {path.c_str('val->')[:-3]}, {cnt_str});\n"
                 else:
-                    ret += f"{'\t'*indent_lvl()}MARSHAL_BYTES_ZEROCOPY(ctx, {path.c_str('val->')[:-3]}, {cnt_path.c_str('val->')});\n"
+                    ret += f"{'\t'*indent_lvl()}MARSHAL_BYTES_ZEROCOPY(ctx, {path.c_str('val->')[:-3]}, {cnt_str});\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"
+            ret += f"{'\t'*indent_lvl()}for ({cnt_typ} {loopvar} = 0; {loopvar} < {cnt_str}; {loopvar}++) {{\n"
             indent_stack.append(IndentLevel(ifdef=False))
         if not isinstance(child.typ, idl.Struct):
             if child.val:
diff --git a/lib9p/protogen/c_unmarshal.py b/lib9p/protogen/c_unmarshal.py
index 018d750..34635f9 100644
--- a/lib9p/protogen/c_unmarshal.py
+++ b/lib9p/protogen/c_unmarshal.py
@@ -93,15 +93,20 @@ def gen_c_unmarshal(versions: set[str], typs: list[idl.UserType]) -> str:
                 )
                 indent_stack.append(IndentLevel(ifdef=True))
         if child.cnt:
-            cnt_path = path.parent().add(child.cnt)
+            if isinstance(child.cnt, int):
+                cnt_str = str(child.cnt)
+                cnt_typ = "size_t"
+            else:
+                cnt_str = path.parent().add(child.cnt).c_str("out->")
+                cnt_typ = c9util.typename(child.cnt.typ)
             if child.typ.static_size == 1:  # SPECIAL (zerocopy)
-                ret += f"{'\t'*indent_lvl()}UNMARSHAL_BYTES(ctx, {path.c_str('out->')[:-3]}, {cnt_path.c_str('out->')});\n"
+                ret += f"{'\t'*indent_lvl()}UNMARSHAL_BYTES(ctx, {path.c_str('out->')[:-3]}, {cnt_str});\n"
                 return idlutil.WalkCmd.KEEP_GOING, pop
             ret += f"{'\t'*indent_lvl()}{path.c_str('out->')[:-3]} = extra;\n"
-            ret += f"{'\t'*indent_lvl()}extra += sizeof({path.c_str('out->')[:-3]}[0]) * {cnt_path.c_str('out->')};\n"
+            ret += f"{'\t'*indent_lvl()}extra += sizeof({path.c_str('out->')[:-3]}[0]) * {cnt_str};\n"
             loopdepth = sum(1 for elem in path.elems if elem.cnt)
             loopvar = chr(ord("i") + loopdepth - 1)
-            ret += f"{'\t'*indent_lvl()}for ({c9util.typename(child.cnt.typ)} {loopvar} = 0; {loopvar} < {cnt_path.c_str('out->')}; {loopvar}++) {{\n"
+            ret += f"{'\t'*indent_lvl()}for ({cnt_typ} {loopvar} = 0; {loopvar} < {cnt_str}; {loopvar}++) {{\n"
             indent_stack.append(IndentLevel(ifdef=False))
         if not isinstance(child.typ, idl.Struct):
             if child.val:
diff --git a/lib9p/protogen/c_validate.py b/lib9p/protogen/c_validate.py
index e315b60..7d0c69e 100644
--- a/lib9p/protogen/c_validate.py
+++ b/lib9p/protogen/c_validate.py
@@ -132,21 +132,33 @@ def gen_c_validate(versions: set[str], typs: list[idl.UserType]) -> str:
         if should_save_offset(parent, child):
             ret += f"{'\t'*indent_lvl()}uint32_t offsetof{''.join('_'+m.membname for m in path.elems)} = net_offset + {incr_buf};\n"
         if child.cnt:
-            assert child.cnt.typ.static_size
-            cnt_path = path.parent().add(child.cnt)
-            incr_flush()
+            if isinstance(child.cnt, int):
+                cnt_str = str(child.cnt)
+                cnt_typ = "size_t"
+            else:
+                assert child.cnt.typ.static_size
+                incr_flush()
+                cnt_str = f"LAST_U{child.cnt.typ.static_size*8}LE()"
+                cnt_typ = c9util.typename(child.cnt.typ)
             if child.membname == "utf8":  # SPECIAL (string)
+                assert child.typ.static_size == 1
                 # Yes, this is content-validation and "belongs" in
                 # gen_validate_content(), not here.  But it's just
                 # easier this way.
-                ret += f"{'\t'*indent_lvl()}VALIDATE_NET_UTF8(LAST_U{child.cnt.typ.static_size*8}LE());\n"
+                incr_flush()
+                ret += f"{'\t'*indent_lvl()}VALIDATE_NET_UTF8({cnt_str});\n"
                 return
             if child.typ.static_size == 1:  # SPECIAL (zerocopy)
-                ret += f"{'\t'*indent_lvl()}VALIDATE_NET_BYTES(LAST_U{child.cnt.typ.static_size*8}LE());\n"
+                if isinstance(child.cnt, int):
+                    incr_buf += child.cnt
+                    return
+                incr_flush()
+                ret += f"{'\t'*indent_lvl()}VALIDATE_NET_BYTES({cnt_str});\n"
                 return
             loopdepth = sum(1 for elem in path.elems if elem.cnt)
             loopvar = chr(ord("i") + loopdepth - 1)
-            ret += f"{'\t'*indent_lvl()}for ({c9util.typename(child.cnt.typ)} {loopvar} = 0, cnt = LAST_U{child.cnt.typ.static_size*8}LE(); {loopvar} < cnt; {loopvar}++) {{\n"
+            incr_flush()
+            ret += f"{'\t'*indent_lvl()}for ({cnt_typ} {loopvar} = 0, cnt = {cnt_str}; {loopvar} < cnt; {loopvar}++) {{\n"
             indent_stack.append(IndentLevel(ifdef=False))
             ret += f"{'\t'*indent_lvl()}RESERVE_HOST_BYTES(sizeof({c9util.typename(child.typ)}));\n"
         if not isinstance(child.typ, idl.Struct):