refactor: improve parsing and testing

assembly/src/ast/mod.rs

@@ -35,7 +35,9 @@ pub use invocation_target::InvocationTarget;
 mod parsers;
 use parsers::{parse_constants, ParserContext};
 
-pub(crate) use parsers::{NAMESPACE_LABEL_PARSER, PROCEDURE_LABEL_PARSER};
+pub(crate) use parsers::{
+    parse_param_with_constant_lookup, NAMESPACE_LABEL_PARSER, PROCEDURE_LABEL_PARSER,
+};
 
 mod serde;
 pub use serde::AstSerdeOptions;

assembly/src/ast/parsers/mod.rs

@@ -122,7 +122,7 @@ fn parse_const_value(
 
 /// Parses a param from the op token with the specified type and index. If the param is a constant
 /// label, it will be looked up in the provided constant map.
-fn parse_param_with_constant_lookup<R>(
+pub(crate) fn parse_param_with_constant_lookup<R>(
     op: &Token,
     param_idx: usize,
     constants: &LocalConstMap,

assembly/src/ast/tests.rs

@@ -960,14 +960,27 @@ fn test_repeat_with_constant_count() {
     begin
         repeat.A
             push.1
-            push.0.1
         end
 
         repeat.B
             push.0
         end
     end";
+
     assert_correct_program_serialization(source, false);
+
+    let nodes: Vec<Node> = vec![
+        Node::Repeat {
+            times: 3,
+            body: CodeBody::new(vec![Node::Instruction(Instruction::PushU8(1))]),
+        },
+        Node::Repeat {
+            times: 14,
+            body: CodeBody::new(vec![Node::Instruction(Instruction::PushU8(0))]),
+        },
+    ];
+
+    assert_program_output(source, BTreeMap::new(), nodes);
 }
 
 fn assert_program_output(source: &str, procedures: LocalProcMap, body: Vec<Node>) {

assembly/src/tests.rs

@@ -1673,7 +1673,7 @@ fn invalid_repeat() {
     if let Err(error) = program {
         assert_eq!(
             error.to_string(),
-            "malformed constant `repeat.23x3` - invalid value: `23x3` - reason: constant with name 23x3 was not initialized"
+            "malformed instruction `repeat.23x3`: parameter '23x3' is invalid"
         );
     }
 }

assembly/src/tokens/mod.rs

@@ -1,6 +1,7 @@
 use super::{
-    ast::InvocationTarget, BTreeMap, ByteReader, ByteWriter, Deserializable, DeserializationError,
-    LibraryPath, ParsingError, ProcedureName, Serializable, String, ToString, Vec,
+    ast::{parse_param_with_constant_lookup, InvocationTarget},
+    BTreeMap, ByteReader, ByteWriter, Deserializable, DeserializationError, LibraryPath,
+    ParsingError, ProcedureName, Serializable, String, ToString, Vec,
 };
 use core::fmt;
 
@@ -243,25 +244,7 @@ impl<'a> Token<'a> {
         match self.num_parts() {
             0 => unreachable!(),
             1 => Err(ParsingError::missing_param(self, "repeat.<num_repetitions>")),
-            2 => {
-                let count = self.parts[1];
-
-                let parsed_number_u64 = if let Ok(number) = count.parse::<u64>() {
-                    number
-                } else {
-                    *constants.get(count).ok_or_else(|| {
-                        ParsingError::invalid_const_value(
-                            self,
-                            count,
-                            &format!("constant with name {} was not initialized", count),
-                        )
-                    })?
-                };
-                let parsed_number: u32 = parsed_number_u64
-                    .try_into()
-                    .map_err(|_| ParsingError::invalid_param(self, 1))?;
-                Ok(parsed_number)
-            }
+            2 => parse_param_with_constant_lookup::<u32>(self, 1, constants),
             _ => Err(ParsingError::extra_param(self)),
         }
     }

docs/src/user_docs/assembly/flow_control.md

@@ -33,7 +33,7 @@ end
 where:
 
 * `instructions` can be a sequence of any instructions, including nested control structures.
-* `count` is the number of times the `instructions` sequence should be repeated (e.g. `repeat.10`). `count` must be an integer greater than $0$.
+* `count` is the number of times the `instructions` sequence should be repeated (e.g. `repeat.10`). `count` must be an integer or a constant greater than $0$.
 
 > **Note**: During compilation the `repeat.<count>` blocks are unrolled and expanded into `<count>` copies of its inner block, there is no additional cost for counting variables in this case.