diff --git a/assembly/src/ast/mod.rs b/assembly/src/ast/mod.rs index 2f1e8c0515..4a6be8bb65 100644 --- a/assembly/src/ast/mod.rs +++ b/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; diff --git a/assembly/src/ast/parsers/mod.rs b/assembly/src/ast/parsers/mod.rs index efa2de6e57..d199514ea5 100644 --- a/assembly/src/ast/parsers/mod.rs +++ b/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( +pub(crate) fn parse_param_with_constant_lookup( op: &Token, param_idx: usize, constants: &LocalConstMap, diff --git a/assembly/src/ast/tests.rs b/assembly/src/ast/tests.rs index 7a51566a89..240040a82e 100644 --- a/assembly/src/ast/tests.rs +++ b/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 = 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) { diff --git a/assembly/src/tests.rs b/assembly/src/tests.rs index 2226125e80..4a015c29b6 100644 --- a/assembly/src/tests.rs +++ b/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" ); } } diff --git a/assembly/src/tokens/mod.rs b/assembly/src/tokens/mod.rs index 25a7102664..fa21be6c89 100644 --- a/assembly/src/tokens/mod.rs +++ b/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.")), - 2 => { - let count = self.parts[1]; - - let parsed_number_u64 = if let Ok(number) = count.parse::() { - 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::(self, 1, constants), _ => Err(ParsingError::extra_param(self)), } } diff --git a/docs/src/user_docs/assembly/flow_control.md b/docs/src/user_docs/assembly/flow_control.md index 8d68bf2ae6..76e7622a70 100644 --- a/docs/src/user_docs/assembly/flow_control.md +++ b/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.` blocks are unrolled and expanded into `` copies of its inner block, there is no additional cost for counting variables in this case.