refactor: emulate immediate value on the parsing stage

0xPolygonMiden · Jun 26, 2024 · 83ddbfc · 83ddbfc
1 parent 32dac16
commit 83ddbfc
Show file tree

Hide file tree

Showing 10 changed files with 48 additions and 78 deletions.
diff --git a/assembly/src/assembler/instruction/mod.rs b/assembly/src/assembler/instruction/mod.rs
@@ -184,14 +184,10 @@ impl Assembler {
             Instruction::U32DivModImm(v) => {
                 u32_ops::u32divmod(span_builder, ctx, Some(v.expect_spanned_value()))?
             }
-            Instruction::U32And => u32_ops::u32and(span_builder, None),
-            Instruction::U32AndImm(v) => u32_ops::u32and(span_builder, Some(v.expect_value())),
-            Instruction::U32Or => u32_ops::u32or(span_builder, None),
-            Instruction::U32OrImm(v) => u32_ops::u32or(span_builder, Some(v.expect_value())),
-            Instruction::U32Xor => u32_ops::u32xor(span_builder, None),
-            Instruction::U32XorImm(v) => u32_ops::u32xor(span_builder, Some(v.expect_value())),
-            Instruction::U32Not => u32_ops::u32not(span_builder, None),
-            Instruction::U32NotImm(v) => u32_ops::u32not(span_builder, Some(v.expect_value())),
+            Instruction::U32And => span_builder.push_op(U32and),
+            Instruction::U32Or => span_builder.push_ops([Dup1, Dup1, U32and, Neg, Add, Add]),
+            Instruction::U32Xor => span_builder.push_op(U32xor),
+            Instruction::U32Not => u32_ops::u32not(span_builder),
             Instruction::U32Shl => u32_ops::u32shl(span_builder, None)?,
             Instruction::U32ShlImm(v) => u32_ops::u32shl(span_builder, Some(v.expect_value()))?,
             Instruction::U32Shr => u32_ops::u32shr(span_builder, None)?,

diff --git a/assembly/src/assembler/instruction/u32_ops.rs b/assembly/src/assembler/instruction/u32_ops.rs
@@ -160,38 +160,13 @@ pub fn u32divmod(
 // BITWISE OPERATIONS
 // ================================================================================================
 
-pub fn u32and(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
-    if let Some(imm) = imm {
-        span_builder.push_op(Push(Felt::from(imm)));
-    }
-    span_builder.push_op(U32and);
-}
-
-pub fn u32or(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
-    if let Some(imm) = imm {
-        span_builder.push_op(Push(Felt::from(imm)));
-    }
-    span_builder.push_ops([Dup1, Dup1, U32and, Neg, Add, Add]);
-}
-
-pub fn u32xor(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
-    if let Some(imm) = imm {
-        span_builder.push_op(Push(Felt::from(imm)));
-    }
-    span_builder.push_op(U32xor);
-}
-
 /// Translates u32not assembly instruction to VM operations.
 ///
 /// The reason this method works is because 2^32 -1 provides a bit mask of ones, which after
 /// subtracting the element, flips the bits of the original value to perform a bitwise NOT.
 ///
 /// This takes 5 VM cycles.
-pub fn u32not(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
-    if let Some(imm) = imm {
-        span_builder.push_op(Push(Felt::from(imm)));
-    }
-
+pub fn u32not(span_builder: &mut BasicBlockBuilder) {
     #[rustfmt::skip]
     let ops = [
         // Perform the operation

diff --git a/assembly/src/ast/instruction/deserialize.rs b/assembly/src/ast/instruction/deserialize.rs
@@ -100,13 +100,9 @@ impl Deserializable for Instruction {
             OpCode::U32DivMod => Ok(Self::U32DivMod),
             OpCode::U32DivModImm => Ok(Self::U32DivModImm(source.read_u32()?.into())),
             OpCode::U32And => Ok(Self::U32And),
-            OpCode::U32AndImm => Ok(Self::U32AndImm(source.read_u32()?.into())),
             OpCode::U32Or => Ok(Self::U32Or),
-            OpCode::U32OrImm => Ok(Self::U32OrImm(source.read_u32()?.into())),
             OpCode::U32Xor => Ok(Self::U32Xor),
-            OpCode::U32XorImm => Ok(Self::U32XorImm(source.read_u32()?.into())),
             OpCode::U32Not => Ok(Self::U32Not),
-            OpCode::U32NotImm => Ok(Self::U32NotImm(source.read_u32()?.into())),
             OpCode::U32Shr => Ok(Self::U32Shr),
             OpCode::U32ShrImm => Ok(Self::U32ShrImm(source.read_u8()?.into())),
             OpCode::U32Shl => Ok(Self::U32Shl),

diff --git a/assembly/src/ast/instruction/mod.rs b/assembly/src/ast/instruction/mod.rs
@@ -107,13 +107,9 @@ pub enum Instruction {
     U32DivMod,
     U32DivModImm(ImmU32),
     U32And,
-    U32AndImm(ImmU32),
     U32Or,
-    U32OrImm(ImmU32),
     U32Xor,
-    U32XorImm(ImmU32),
     U32Not,
-    U32NotImm(ImmU32),
     U32Shr,
     U32ShrImm(ImmU8),
     U32Shl,

diff --git a/assembly/src/ast/instruction/opcode.rs b/assembly/src/ast/instruction/opcode.rs
@@ -98,13 +98,9 @@ pub enum OpCode {
     U32DivMod,
     U32DivModImm,
     U32And,
-    U32AndImm,
     U32Or,
-    U32OrImm,
     U32Xor,
-    U32XorImm,
     U32Not,
-    U32NotImm,
     U32Shr,
     U32ShrImm,
     U32Shl,

diff --git a/assembly/src/ast/instruction/print.rs b/assembly/src/ast/instruction/print.rs
@@ -111,13 +111,9 @@ impl PrettyPrint for Instruction {
             Self::U32DivMod => const_text("u32divmod"),
             Self::U32DivModImm(value) => inst_with_imm("u32divmod", value),
             Self::U32And => const_text("u32and"),
-            Self::U32AndImm(value) => inst_with_imm("u32and", value),
             Self::U32Or => const_text("u32or"),
-            Self::U32OrImm(value) => inst_with_imm("u32or", value),
             Self::U32Xor => const_text("u32xor"),
-            Self::U32XorImm(value) => inst_with_imm("u32xor", value),
             Self::U32Not => const_text("u32not"),
-            Self::U32NotImm(value) => inst_with_imm("u32not", value),
             Self::U32Shr => const_text("u32shr"),
             Self::U32ShrImm(value) => inst_with_imm("u32shr", value),
             Self::U32Shl => const_text("u32shl"),

diff --git a/assembly/src/ast/instruction/serialize.rs b/assembly/src/ast/instruction/serialize.rs
@@ -176,25 +176,9 @@ impl Serializable for Instruction {
                 target.write_u32(v.expect_value());
             }
             Self::U32And => OpCode::U32And.write_into(target),
-            Self::U32AndImm(v) => {
-                OpCode::U32AndImm.write_into(target);
-                target.write_u32(v.expect_value());
-            }
             Self::U32Or => OpCode::U32Or.write_into(target),
-            Self::U32OrImm(v) => {
-                OpCode::U32OrImm.write_into(target);
-                target.write_u32(v.expect_value());
-            }
             Self::U32Xor => OpCode::U32Xor.write_into(target),
-            Self::U32XorImm(v) => {
-                OpCode::U32XorImm.write_into(target);
-                target.write_u32(v.expect_value());
-            }
             Self::U32Not => OpCode::U32Not.write_into(target),
-            Self::U32NotImm(v) => {
-                OpCode::U32NotImm.write_into(target);
-                target.write_u32(v.expect_value());
-            }
             Self::U32Shr => OpCode::U32Shr.write_into(target),
             Self::U32ShrImm(v) => {
                 OpCode::U32ShrImm.write_into(target);

diff --git a/assembly/src/ast/visit.rs b/assembly/src/ast/visit.rs
@@ -314,10 +314,6 @@ where
         | U32GteImm(ref imm)
         | U32MinImm(ref imm)
         | U32MaxImm(ref imm)
-        | U32AndImm(ref imm)
-        | U32OrImm(ref imm)
-        | U32XorImm(ref imm)
-        | U32NotImm(ref imm)
         | MemLoadImm(ref imm)
         | MemLoadWImm(ref imm)
         | MemStoreImm(ref imm)
@@ -767,10 +763,6 @@ where
         | U32GteImm(ref mut imm)
         | U32MinImm(ref mut imm)
         | U32MaxImm(ref mut imm)
-        | U32AndImm(ref mut imm)
-        | U32OrImm(ref mut imm)
-        | U32XorImm(ref mut imm)
-        | U32NotImm(ref mut imm)
         | MemLoadImm(ref mut imm)
         | MemLoadWImm(ref mut imm)
         | MemStoreImm(ref mut imm)

diff --git a/assembly/src/parser/grammar.lalrpop b/assembly/src/parser/grammar.lalrpop
@@ -716,30 +716,50 @@ FoldableInstWithU32Immediate: SmallOpsVec = {
         let span = span!(l, r);
         match imm {
             Some(imm) if imm == 0 => smallvec![Op::Inst(Span::new(span, Instruction::Drop)), Op::Inst(Span::new(span, Instruction::PushU8(0)))],
-            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32AndImm(imm)))],
+            Some(imm) => {
+                match imm {
+                    Immediate::Constant(name) => smallvec![Op::Inst(Span::new(span, Instruction::Push(Immediate::Constant(name)))), Op::Inst(Span::new(span, Instruction::U32And))],
+                    Immediate::Value(value) => smallvec![Op::Inst(Span::new(span, Instruction::PushU32(value.into_inner()))), Op::Inst(Span::new(span, Instruction::U32And))],
+                }
+            }
             None => smallvec![Op::Inst(Span::new(span, Instruction::U32And))],
         }
     },
     <l:@L> "u32or" <imm:MaybeImm<U32>> <r:@R> => {
         let span = span!(l, r);
         match imm {
             Some(imm) if imm == 0 => smallvec![],
-            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32OrImm(imm)))],
+            Some(imm) => {
+                match imm {
+                    Immediate::Constant(name) => smallvec![Op::Inst(Span::new(span, Instruction::Push(Immediate::Constant(name)))), Op::Inst(Span::new(span, Instruction::U32Or))],
+                    Immediate::Value(value) => smallvec![Op::Inst(Span::new(span, Instruction::PushU32(value.into_inner()))), Op::Inst(Span::new(span, Instruction::U32Or))],
+                }
+            }
             None => smallvec![Op::Inst(Span::new(span, Instruction::U32Or))],
         }
     },
     <l:@L> "u32xor" <imm:MaybeImm<U32>> <r:@R> => {
         let span = span!(l, r);
         match imm {
             Some(imm) if imm == 0 => smallvec![],
-            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32XorImm(imm)))],
+            Some(imm) => {
+                match imm {
+                    Immediate::Constant(name) => smallvec![Op::Inst(Span::new(span, Instruction::Push(Immediate::Constant(name)))), Op::Inst(Span::new(span, Instruction::U32Xor))],
+                    Immediate::Value(value) => smallvec![Op::Inst(Span::new(span, Instruction::PushU32(value.into_inner()))), Op::Inst(Span::new(span, Instruction::U32Xor))],
+                }
+            }
             None => smallvec![Op::Inst(Span::new(span, Instruction::U32Xor))],
         }
     },
     <l:@L> "u32not" <imm:MaybeImm<U32>> <r:@R> => {
         let span = span!(l, r);
         match imm {
-            Some(imm) => smallvec![Op::Inst(Span::new(span, Instruction::U32NotImm(imm)))],
+            Some(imm) => {
+                match imm {
+                    Immediate::Constant(name) => smallvec![Op::Inst(Span::new(span, Instruction::Push(Immediate::Constant(name)))), Op::Inst(Span::new(span, Instruction::U32Not))],
+                    Immediate::Value(value) => smallvec![Op::Inst(Span::new(span, Instruction::PushU32(value.into_inner()))), Op::Inst(Span::new(span, Instruction::U32Not))],
+                }
+            }
             None => smallvec![Op::Inst(Span::new(span, Instruction::U32Not))],
         }
     },

diff --git a/assembly/src/tests.rs b/assembly/src/tests.rs
@@ -461,6 +461,25 @@ end";
     Ok(())
 }
 
+#[test]
+fn simple_bitwise_constant() -> TestResult {
+    let mut context = TestContext::default();
+    let source = source_file!(
+        "\
+    const.TEST_CONSTANT=7
+    begin
+        push.2 u32and.TEST_CONSTANT
+    end"
+    );
+    let expected = "\
+begin
+    basic_block push(2) push(7) u32and end
+end";
+    let program = context.assemble(source)?;
+    assert_str_eq!(format!("{program}"), expected);
+    Ok(())
+}
+
 #[test]
 fn multiple_constants_push() -> TestResult {
     let mut context = TestContext::default();