fix bug in typechecking named theorems that are not onestate

wilcoxjay · Apr 9, 2024 · ddbd066 · ddbd066
1 parent 29bb75e
commit ddbd066
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/typechecker.py b/src/typechecker.py
@@ -434,11 +434,16 @@ def typecheck_tracedecl(scope: syntax.Scope, d: syntax.TraceDecl) -> None:
             assert False
 
 
-def add_named_macro(prog: syntax.Program, name: Optional[str], expr: syntax.Expr) -> None:
+def add_named_macro(
+        prog: syntax.Program,
+        name: Optional[str],
+        expr: syntax.Expr,
+        num_states: int = 1,
+) -> None:
     if name is not None:
         prog.decls.append(DefinitionDecl(
             is_public_transition=False,
-            num_states=1,
+            num_states=num_states,
             name=name,
             params=(),
             mods=(),
@@ -465,7 +470,7 @@ def typecheck_program_vocab(prog: syntax.Program) -> None:
         add_named_macro(prog, a.name, a.expr)
 
     for t in prog.theorems():
-        add_named_macro(prog, t.name, t.expr)
+        add_named_macro(prog, t.name, t.expr, num_states=t.num_states)
 
     add_named_macro(prog, 'safety', syntax.And(*(s.expr for s in prog.safeties())))