addressed leo's comments'

qwang98 · Aug 17, 2023 · 1be48e1 · 1be48e1
1 parent 06c898c
commit 1be48e1
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/python/chiquito/dsl.py b/python/chiquito/dsl.py
@@ -83,21 +83,19 @@ def pragma_disable_q_enable(self: Circuit) -> None:
 
     def add(self: Circuit, step_type: StepType, args: Any):
         assert self.mode == CircuitMode.Trace
-        if self.num_step_instances >= self.ast.num_steps:
+        if len(self.witness.step_instances) >= self.ast.num_steps:
             raise ValueError(f"Number of step instances exceeds {self.ast.num_steps}")
-        self.num_step_instances += 1
         step_instance: StepInstance = step_type.gen_step_instance(args)
         self.witness.step_instances.append(step_instance)
 
     def needs_padding(self: Circuit) -> bool:
-        return self.num_step_instances < self.ast.num_steps
+        return len(self.witness.step_instances) < self.ast.num_steps
 
     def padding(self: Circuit, step_type: StepType, args: Any):
         while self.needs_padding():
             self.add(step_type, args)
 
     def gen_witness(self: Circuit, args: Any) -> TraceWitness:
-        self.num_step_instances = 0
         self.mode = CircuitMode.Trace
         self.witness = TraceWitness()
         self.trace(args)