Merge pull request #1124 from 0xPolygonMiden/andrew-support-constants…

…-on-control-flow

Support constants in MASM `repeat` loops

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Expanded capabilities of the `debug` decorator. Added `debug.mem` and `debug.local` variations (#1103).
 - Introduced the `emit.<event_id>` assembly instruction (#1119).
 - Introduced the `procref.<proc_name>` assembly instruction (#1113).
+- Added the ability to use constants as counters in `repeat` loops (#1124). 
 
 #### Stdlib
 - Introduced `std::utils` module with `is_empty_word` procedure.  Refactored `std::collections::smt`

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/context.rs

@@ -133,7 +133,7 @@ impl ParserContext<'_> {
         // record start of the repeat block and consume the 'repeat' token
         let repeat_start = tokens.pos();
         let repeat_token = tokens.read().expect("no repeat token");
-        let times = repeat_token.parse_repeat()?;
+        let times = repeat_token.parse_repeat(&self.local_constants)?;
         tokens.advance();
 
         // read the loop body

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

@@ -951,6 +951,38 @@ fn test_ast_module_serde_imports_not_serialized() {
     assert_correct_module_serialization(source, false);
 }
 
+#[test]
+fn test_repeat_with_constant_count() {
+    let source = "\
+    const.A=3
+    const.B=A*3+5
+
+    begin
+        repeat.A
+            push.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>) {
     let program = ProgramAst::parse(source).unwrap();
     assert_eq!(program.body.nodes(), body);

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;
 
@@ -238,12 +239,12 @@ impl<'a> Token<'a> {
         }
     }
 
-    pub fn parse_repeat(&self) -> Result<u32, ParsingError> {
+    pub fn parse_repeat(&self, constants: &BTreeMap<String, u64>) -> Result<u32, ParsingError> {
         assert_eq!(Self::REPEAT, self.parts[0], "not a repeat");
         match self.num_parts() {
             0 => unreachable!(),
             1 => Err(ParsingError::missing_param(self, "repeat.<num_repetitions>")),
-            2 => self.parts[1].parse::<u32>().map_err(|_| ParsingError::invalid_param(self, 1)),
+            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](./code_organization.md#constants) 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.