feat: add support for imm values for u32 comparison instructions (#1358)

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added support for immediate values for `lt`, `lte`, `gt`, `gte` comparison instructions (#1346).
 - Change MAST to a table-based representation (#1349)
 - Adjusted prover's metal acceleration code to work with 0.9 versions of the crates (#1357)
+- Added support for immediate values for `u32lt`, `u32lte`, `u32gt`, `u32gte`, `u32min` and `u32max` comparison instructions (#1358).
 
 ## 0.9.2 (2024-05-22) - `stdlib` crate only
 - Skip writing MASM documentation to file when building on docs.rs (#1341).

assembly/src/assembler/instruction/mod.rs

@@ -200,12 +200,18 @@ impl Assembler {
             Instruction::U32Ctz => u32_ops::u32ctz(span_builder),
             Instruction::U32Clo => u32_ops::u32clo(span_builder),
             Instruction::U32Cto => u32_ops::u32cto(span_builder),
-            Instruction::U32Lt => u32_ops::u32lt(span_builder),
-            Instruction::U32Lte => u32_ops::u32lte(span_builder),
-            Instruction::U32Gt => u32_ops::u32gt(span_builder),
-            Instruction::U32Gte => u32_ops::u32gte(span_builder),
-            Instruction::U32Min => u32_ops::u32min(span_builder),
-            Instruction::U32Max => u32_ops::u32max(span_builder),
+            Instruction::U32Lt => u32_ops::u32lt(span_builder, None),
+            Instruction::U32LtImm(v) => u32_ops::u32lt(span_builder, Some(v.expect_value())),
+            Instruction::U32Lte => u32_ops::u32lte(span_builder, None),
+            Instruction::U32LteImm(v) => u32_ops::u32lte(span_builder, Some(v.expect_value())),
+            Instruction::U32Gt => u32_ops::u32gt(span_builder, None),
+            Instruction::U32GtImm(v) => u32_ops::u32gt(span_builder, Some(v.expect_value())),
+            Instruction::U32Gte => u32_ops::u32gte(span_builder, None),
+            Instruction::U32GteImm(v) => u32_ops::u32gte(span_builder, Some(v.expect_value())),
+            Instruction::U32Min => u32_ops::u32min(span_builder, None),
+            Instruction::U32MinImm(v) => u32_ops::u32min(span_builder, Some(v.expect_value())),
+            Instruction::U32Max => u32_ops::u32max(span_builder, None),
+            Instruction::U32MaxImm(v) => u32_ops::u32max(span_builder, Some(v.expect_value())),
 
             // ----- stack manipulation -----------------------------------------------------------
             Instruction::Drop => span_builder.push_op(Drop),

assembly/src/assembler/instruction/u32_ops.rs

@@ -714,17 +714,29 @@ fn calculate_cto(span: &mut BasicBlockBuilder) {
 // COMPARISON OPERATIONS
 // ================================================================================================
 
-/// Translates u32lt assembly instructions to VM operations.
+/// Translates u32lt assembly instruction to VM operations.
 ///
-/// This operation takes 3 cycles.
-pub fn u32lt(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 3 cycles without immediate value.
+/// - 4 cycles with immediate value.
+pub fn u32lt(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     compute_lt(span_builder);
 }
 
-/// Translates u32lte assembly instructions to VM operations.
+/// Translates u32lte assembly instruction to VM operations.
 ///
-/// This operation takes 5 cycles.
-pub fn u32lte(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 5 cycles without immediate value.
+/// - 6 cycles with immediate value.
+pub fn u32lte(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     // Compute the lt with reversed number to get a gt check
     span_builder.push_op(Swap);
     compute_lt(span_builder);
@@ -733,48 +745,72 @@ pub fn u32lte(span_builder: &mut BasicBlockBuilder) {
     span_builder.push_op(Not);
 }
 
-/// Translates u32gt assembly instructions to VM operations.
+/// Translates u32gt assembly instruction to VM operations.
 ///
-/// This operation takes 4 cycles.
-pub fn u32gt(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 4 cycles without immediate value.
+/// - 5 cycles with immediate value.
+pub fn u32gt(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     // Reverse the numbers so we can get a gt check.
     span_builder.push_op(Swap);
 
     compute_lt(span_builder);
 }
 
-/// Translates u32gte assembly instructions to VM operations.
+/// Translates u32gte assembly instruction to VM operations.
 ///
-/// This operation takes 4 cycles.
-pub fn u32gte(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 4 cycles without immediate value.
+/// - 5 cycles with immediate value.
+pub fn u32gte(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     compute_lt(span_builder);
 
     // Flip the final results to get the gte results.
     span_builder.push_op(Not);
 }
 
-/// Translates u32min assembly instructions to VM operations.
+/// Translates u32min assembly instruction to VM operations.
 ///
 /// Specifically, we subtract the top value from the second to the top value (U32SUB), check the
 /// underflow flag (EQZ), and perform a conditional swap (CSWAP) to have the max number in front.
 /// Then we finally drop the top element to keep the min.
 ///
-/// This operation takes 8 cycles.
-pub fn u32min(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 8 cycles without immediate value.
+/// - 9 cycles with immediate value.
+pub fn u32min(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     compute_max_and_min(span_builder);
 
     // Drop the max and keep the min
     span_builder.push_op(Drop);
 }
 
-/// Translates u32max assembly instructions to VM operations.
+/// Translates u32max assembly instruction to VM operations.
 ///
 /// Specifically, we subtract the top value from the second to the top value (U32SUB), check the
 /// underflow flag (EQZ), and perform a conditional swap (CSWAP) to have the max number in front.
 /// Then we finally drop the 2nd element to keep the max.
 ///
-/// This operation takes 9 cycles.
-pub fn u32max(span_builder: &mut BasicBlockBuilder) {
+/// This operation takes:
+/// - 9 cycles without immediate value.
+/// - 10 cycles with immediate value.
+pub fn u32max(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     compute_max_and_min(span_builder);
 
     // Drop the min and keep the max

assembly/src/ast/instruction/deserialize.rs

@@ -116,11 +116,17 @@ impl Deserializable for Instruction {
             OpCode::U32Clo => Ok(Self::U32Clo),
             OpCode::U32Cto => Ok(Self::U32Cto),
             OpCode::U32Lt => Ok(Self::U32Lt),
+            OpCode::U32LtImm => Ok(Self::U32LtImm(source.read_u32()?.into())),
             OpCode::U32Lte => Ok(Self::U32Lte),
+            OpCode::U32LteImm => Ok(Self::U32LteImm(source.read_u32()?.into())),
             OpCode::U32Gt => Ok(Self::U32Gt),
+            OpCode::U32GtImm => Ok(Self::U32GtImm(source.read_u32()?.into())),
             OpCode::U32Gte => Ok(Self::U32Gte),
+            OpCode::U32GteImm => Ok(Self::U32GteImm(source.read_u32()?.into())),
             OpCode::U32Min => Ok(Self::U32Min),
+            OpCode::U32MinImm => Ok(Self::U32MinImm(source.read_u32()?.into())),
             OpCode::U32Max => Ok(Self::U32Max),
+            OpCode::U32MaxImm => Ok(Self::U32MaxImm(source.read_u32()?.into())),
 
             // ----- stack manipulation -----------------------------------------------------------
             OpCode::Drop => Ok(Self::Drop),

assembly/src/ast/instruction/mod.rs

@@ -123,11 +123,17 @@ pub enum Instruction {
     U32Clo,
     U32Cto,
     U32Lt,
+    U32LtImm(ImmU32),
     U32Lte,
+    U32LteImm(ImmU32),
     U32Gt,
+    U32GtImm(ImmU32),
     U32Gte,
+    U32GteImm(ImmU32),
     U32Min,
+    U32MinImm(ImmU32),
     U32Max,
+    U32MaxImm(ImmU32),
 
     // ----- stack manipulation ------------------------------------------------------------------
     Drop,

assembly/src/ast/instruction/opcode.rs

@@ -114,11 +114,17 @@ pub enum OpCode {
     U32Clo,
     U32Cto,
     U32Lt,
+    U32LtImm,
     U32Lte,
+    U32LteImm,
     U32Gt,
+    U32GtImm,
     U32Gte,
+    U32GteImm,
     U32Min,
+    U32MinImm,
     U32Max,
+    U32MaxImm,
 
     // ----- stack manipulation ------------------------------------------------------------------
     Drop,

assembly/src/ast/instruction/print.rs

@@ -127,11 +127,17 @@ impl PrettyPrint for Instruction {
             Self::U32Clo => const_text("u32clo"),
             Self::U32Cto => const_text("u32cto"),
             Self::U32Lt => const_text("u32lt"),
+            Self::U32LtImm(value) => inst_with_imm("u32lt", value),
             Self::U32Lte => const_text("u32lte"),
+            Self::U32LteImm(value) => inst_with_imm("u32lte", value),
             Self::U32Gt => const_text("u32gt"),
+            Self::U32GtImm(value) => inst_with_imm("u32gt", value),
             Self::U32Gte => const_text("u32gte"),
+            Self::U32GteImm(value) => inst_with_imm("u32gte", value),
             Self::U32Min => const_text("u32min"),
+            Self::U32MinImm(value) => inst_with_imm("u32min", value),
             Self::U32Max => const_text("u32max"),
+            Self::U32MaxImm(value) => inst_with_imm("u32min", value),
 
             // ----- stack manipulation -----------------------------------------------------------
             Self::Drop => const_text("drop"),

assembly/src/ast/instruction/serialize.rs

@@ -204,11 +204,35 @@ impl Serializable for Instruction {
             Self::U32Clo => OpCode::U32Clo.write_into(target),
             Self::U32Cto => OpCode::U32Cto.write_into(target),
             Self::U32Lt => OpCode::U32Lt.write_into(target),
+            Self::U32LtImm(v) => {
+                OpCode::U32LtImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
             Self::U32Lte => OpCode::U32Lte.write_into(target),
+            Self::U32LteImm(v) => {
+                OpCode::U32LteImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
             Self::U32Gt => OpCode::U32Gt.write_into(target),
+            Self::U32GtImm(v) => {
+                OpCode::U32GtImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
             Self::U32Gte => OpCode::U32Gte.write_into(target),
+            Self::U32GteImm(v) => {
+                OpCode::U32GteImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
             Self::U32Min => OpCode::U32Min.write_into(target),
+            Self::U32MinImm(v) => {
+                OpCode::U32MinImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
             Self::U32Max => OpCode::U32Max.write_into(target),
+            Self::U32MaxImm(v) => {
+                OpCode::U32MaxImm.write_into(target);
+                target.write_u32(v.expect_value());
+            }
 
             // ----- stack manipulation -----------------------------------------------------------
             Self::Drop => OpCode::Drop.write_into(target),

assembly/src/ast/visit.rs

@@ -308,6 +308,12 @@ where
         | U32DivImm(ref imm)
         | U32ModImm(ref imm)
         | U32DivModImm(ref imm)
+        | U32LtImm(ref imm)
+        | U32LteImm(ref imm)
+        | U32GtImm(ref imm)
+        | U32GteImm(ref imm)
+        | U32MinImm(ref imm)
+        | U32MaxImm(ref imm)
         | MemLoadImm(ref imm)
         | MemLoadWImm(ref imm)
         | MemStoreImm(ref imm)
@@ -751,6 +757,12 @@ where
         | U32DivImm(ref mut imm)
         | U32ModImm(ref mut imm)
         | U32DivModImm(ref mut imm)
+        | U32LtImm(ref mut imm)
+        | U32LteImm(ref mut imm)
+        | U32GtImm(ref mut imm)
+        | U32GteImm(ref mut imm)
+        | U32MinImm(ref mut imm)
+        | U32MaxImm(ref mut imm)
         | MemLoadImm(ref mut imm)
         | MemLoadWImm(ref mut imm)
         | MemStoreImm(ref mut imm)

assembly/src/parser/grammar.lalrpop

@@ -447,12 +447,6 @@ Inst: Instruction = {
     "swapdw" => Instruction::SwapDw,
     "u32and" => Instruction::U32And,
     "u32cast" => Instruction::U32Cast,
-    "u32gt" => Instruction::U32Gt,
-    "u32gte" => Instruction::U32Gte,
-    "u32lt" => Instruction::U32Lt,
-    "u32lte" => Instruction::U32Lte,
-    "u32max" => Instruction::U32Max,
-    "u32min" => Instruction::U32Min,
     "u32not" => Instruction::U32Not,
     "u32or" => Instruction::U32Or,
     "u32overflowing_add3" => Instruction::U32OverflowingAdd3,
@@ -764,6 +758,48 @@ FoldableInstWithU32Immediate: SmallOpsVec = {
             None => smallvec![Op::Inst(Span::new(span, Instruction::U32Rotr))],
         }
     },
+    <l:@L> "u32lt" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32LtImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Lt))],
+        }
+    },
+    <l:@L> "u32lte" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32LteImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Lte))],
+        }
+    },
+    <l:@L> "u32gt" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32GtImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Gt))],
+        }
+    },
+    <l:@L> "u32gte" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32GteImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Gte))],
+        }
+    },
+    <l:@L> "u32min" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32MinImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Min))],
+        }
+    },
+    <l:@L> "u32max" <imm:MaybeImm<U32>> <r:@R> => {
+        let span = span!(l, r);
+        match imm {
+            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32MaxImm(imm)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Max))],
+        }
+    },
 }
 
 #[inline]
@@ -976,7 +1012,6 @@ Shift32: u8 = {
     }
 }
 
-
 #[inline]
 RawU8: u8 = {
     <n:U8> => n.into_inner(),

docs/src/user_docs/assembly/u32_operations.md

@@ -65,9 +65,9 @@ If the error code is omitted, the default value of $0$ is assumed.
 
 | Instruction                                                                      | Stack input  | Stack output    | Notes                                                                                                                                                                                                                  |
 | -------------------------------------------------------------------------------- | ------------ | --------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| u32lt <br> - *(3 cycles)*                                              | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a < b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
-| u32lte <br> - *(5 cycles)*                                             | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a \le b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                    |
-| u32gt <br> - *(4 cycles)*                                              | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a > b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
-| u32gte <br> - *(4 cycles)*                                             | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a \ge b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                    |
-| u32min <br> - *(8 cycles)*                                             | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} a, & \text{if}\ a < b \\ b, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
-| u32max <br> - *(9 cycles)*                                             | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} a, & \text{if}\ a > b \\ b, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
+| u32lt <br> - *(3 cycles)* <br> u32lt.*b* <br> - *(4 cycles)*           | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a < b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
+| u32lte <br> - *(5 cycles)* <br> u32lte.*b* <br> - *(6 cycles)*         | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a \le b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                    |
+| u32gt <br> - *(4 cycles)* <br> u32gt.*b* <br> - *(5 cycles)*           | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a > b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
+| u32gte <br> - *(4 cycles)* <br> u32gte.*b* <br> - *(5 cycles)*         | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} 1, & \text{if}\ a \ge b \\ 0, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                    |
+| u32min <br> - *(8 cycles)* <br> u32min.*b* <br> - *(9 cycles)*         | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} a, & \text{if}\ a < b \\ b, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |
+| u32max <br> - *(9 cycles)* <br> u32max.*b* <br> - *(10 cycles)*        | [b, a, ...]  | [c, ...]        | $c \leftarrow \begin{cases} a, & \text{if}\ a > b \\ b, & \text{otherwise}\ \end{cases}$ <br> Undefined if $max(a, b) \ge 2^{32}$                                                                                      |