feat: implement immediate values for u32 binary instructions,

implement binary parser for lexer

assembly/src/assembler/instruction/mod.rs

@@ -183,10 +183,14 @@ impl Assembler {
             Instruction::U32DivModImm(v) => {
                 u32_ops::u32divmod(span_builder, ctx, Some(v.expect_spanned_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::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::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)?,

assembly/src/assembler/instruction/u32_ops.rs

@@ -160,13 +160,38 @@ 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) {
+pub fn u32not(span_builder: &mut BasicBlockBuilder, imm: Option<u32>) {
+    if let Some(imm) = imm {
+        span_builder.push_op(Push(Felt::from(imm)));
+    }
+
     #[rustfmt::skip]
     let ops = [
         // Perform the operation

assembly/src/ast/instruction/deserialize.rs

@@ -99,9 +99,13 @@ 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),

assembly/src/ast/instruction/mod.rs

@@ -106,9 +106,13 @@ pub enum Instruction {
     U32DivMod,
     U32DivModImm(ImmU32),
     U32And,
+    U32AndImm(ImmU32),
     U32Or,
+    U32OrImm(ImmU32),
     U32Xor,
+    U32XorImm(ImmU32),
     U32Not,
+    U32NotImm(ImmU32),
     U32Shr,
     U32ShrImm(ImmU8),
     U32Shl,

assembly/src/ast/instruction/opcode.rs

@@ -97,9 +97,13 @@ pub enum OpCode {
     U32DivMod,
     U32DivModImm,
     U32And,
+    U32AndImm,
     U32Or,
+    U32OrImm,
     U32Xor,
+    U32XorImm,
     U32Not,
+    U32NotImm,
     U32Shr,
     U32ShrImm,
     U32Shl,

assembly/src/ast/instruction/print.rs

@@ -110,9 +110,13 @@ 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"),

assembly/src/ast/instruction/serialize.rs

@@ -175,9 +175,25 @@ 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);

assembly/src/ast/visit.rs

@@ -308,6 +308,10 @@ where
         | U32DivImm(ref imm)
         | U32ModImm(ref imm)
         | U32DivModImm(ref imm)
+        | U32AndImm(ref imm)
+        | U32OrImm(ref imm)
+        | U32XorImm(ref imm)
+        | U32NotImm(ref imm)
         | MemLoadImm(ref imm)
         | MemLoadWImm(ref imm)
         | MemStoreImm(ref imm)
@@ -751,6 +755,10 @@ where
         | U32DivImm(ref mut imm)
         | U32ModImm(ref mut imm)
         | U32DivModImm(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)

assembly/src/parser/error.rs

@@ -162,6 +162,12 @@ pub enum ParsingError {
         span: SourceSpan,
         kind: HexErrorKind,
     },
+    #[error("Binary value overflowed the field modulus")]
+    #[diagnostic()]
+    InvalidBinaryLiteral {
+        #[label]
+        span: SourceSpan,
+    },
     #[error("invalid MAST root literal")]
     InvalidMastRoot {
         #[label]
@@ -335,6 +341,7 @@ fn simplify_expected_tokens(expected: Vec<String>) -> Vec<String> {
                 "quoted_ident" => return Some("quoted identifier".to_string()),
                 "doc_comment" => return Some("doc comment".to_string()),
                 "hex_value" => return Some("hex-encoded literal".to_string()),
+                "bin_value" => return Some("bin-encoded literal".to_string()),
                 "uint" => return Some("integer literal".to_string()),
                 "EOF" => return Some("end of file".to_string()),
                 other => other[1..].strip_suffix('"').and_then(|t| Token::parse(t).ok()),

assembly/src/parser/grammar.lalrpop

@@ -12,7 +12,7 @@ use vm_core::{Felt, FieldElement, StarkField, crypto::hash::RpoDigest};
 
 use crate::{LibraryPath, LibraryNamespace, ast::*, diagnostics::SourceFile, SourceSpan};
 use super::{
-    HexEncodedValue, Token, ParseError, ParsingError,
+    BinEncodedValue, HexEncodedValue, Token, ParseError, ParsingError,
     LiteralErrorKind, HexErrorKind, Span, Spanned,
     DocumentationType
 };
@@ -35,6 +35,7 @@ extern {
         const_ident => Token::ConstantIdent(<&'input str>),
         quoted_ident => Token::QuotedIdent(<&'input str>),
         hex_value => Token::HexValue(<HexEncodedValue>),
+        bin_value => Token::BinValue(<BinEncodedValue>),
         doc_comment => Token::DocComment(<DocumentationType>),
         comment => Token::Comment,
         uint => Token::Int(<u64>),
@@ -445,16 +446,13 @@ Inst: Instruction = {
     "rcomb_base" => Instruction::RCombBase,
     "sdepth" => Instruction::Sdepth,
     "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,
     "u32overflowing_madd" => Instruction::U32OverflowingMadd,
     "u32popcnt" => Instruction::U32Popcnt,
@@ -467,7 +465,6 @@ Inst: Instruction = {
     "u32testw" => Instruction::U32TestW,
     "u32wrapping_add3" => Instruction::U32WrappingAdd3,
     "u32wrapping_madd" => Instruction::U32WrappingMadd,
-    "u32xor" => Instruction::U32Xor,
     "xor" => Instruction::Xor,
 }
 
@@ -686,6 +683,37 @@ FoldableInstWithU32Immediate: SmallOpsVec = {
             None => Ok(smallvec![Op::Inst(Span::new(span, Instruction::U32Mod))]),
         }
     },
+    <l:@L> "u32and" <imm:MaybeImm<U32>> <r:@R> => {
+        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)))],
+            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)))],
+            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)))],
+            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)))],
+            None => smallvec![Op::Inst(Span::new(span, Instruction::U32Not))],
+        }
+    },
     <l:@L> "u32wrapping_add" <imm:MaybeImm<U32>> <r:@R> => {
         let span = span!(l, r);
         match imm {
@@ -1013,6 +1041,15 @@ U32: u32 = {
             HexEncodedValue::U32(v) => Ok(v),
             _ =>  Err(ParseError::User { error: ParsingError::InvalidLiteral { span: span!(l, r), kind: LiteralErrorKind::U32Overflow } }),
         }
+    },
+
+    <l:@L> <value:bin_value> <r:@R> =>? {
+        match value {
+            BinEncodedValue::U8(v) => Ok(v as u32),
+            BinEncodedValue::U16(v) => Ok(v as u32),
+            BinEncodedValue::U32(v) => Ok(v),
+            _ =>  Err(ParseError::User { error: ParsingError::InvalidLiteral { span: span!(l, r), kind: LiteralErrorKind::U32Overflow } }),
+        }
     }
 }
 

assembly/src/parser/lexer.rs

@@ -3,8 +3,8 @@ use alloc::string::String;
 use core::{num::IntErrorKind, ops::Range};
 
 use super::{
-    DocumentationType, HexEncodedValue, HexErrorKind, LiteralErrorKind, ParsingError, Scanner,
-    SourceSpan, Token,
+    BinEncodedValue, DocumentationType, HexEncodedValue, HexErrorKind, LiteralErrorKind,
+    ParsingError, Scanner, SourceSpan, Token,
 };
 
 /// The value produced by the [Lexer] when iterated
@@ -293,6 +293,11 @@ impl<'input> Lexer<'input> {
                     self.skip();
                     self.lex_hex()
                 }
+                'b' => {
+                    self.skip();
+                    self.skip();
+                    self.lex_bin()
+                }
                 '0'..='9' => self.lex_number(),
                 _ => pop!(self, Token::Int(0)),
             },
@@ -524,6 +529,28 @@ impl<'input> Lexer<'input> {
         let value = parse_hex(span, self.slice_span(digit_start..end))?;
         Ok(Token::HexValue(value))
     }
+
+    fn lex_bin(&mut self) -> Result<Token<'input>, ParsingError> {
+        // Expect the first character to be a valid binary digit
+        debug_assert!(is_ascii_binary(self.read()));
+
+        loop {
+            // If we hit a non-binary digit, we're done
+            let c1 = self.read();
+            if !c1.is_ascii_hexdigit() {
+                break;
+            }
+            self.skip();
+        }
+
+        let span = self.span();
+        let start = span.start() as u32;
+        let digit_start = start + 2;
+        let end = span.end() as u32;
+        let span = SourceSpan::from(start..end);
+        let value = parse_bin(span, self.slice_span(digit_start..end))?;
+        Ok(Token::BinValue(value))
+    }
 }
 
 impl<'input> Iterator for Lexer<'input> {
@@ -561,7 +588,7 @@ fn parse_hex(span: SourceSpan, hex_digits: &str) -> Result<HexEncodedValue, Pars
                     kind: LiteralErrorKind::FeltOverflow,
                 });
             }
-            Ok(shrink_u64(value))
+            Ok(shrink_u64_hex(value))
         }
         // Word
         64 => {
@@ -609,8 +636,35 @@ fn parse_hex(span: SourceSpan, hex_digits: &str) -> Result<HexEncodedValue, Pars
     }
 }
 
+fn parse_bin(span: SourceSpan, bin_digits: &str) -> Result<BinEncodedValue, ParsingError> {
+    use vm_core::StarkField;
+    if bin_digits.len() <= 64 {
+        let value =
+            u64::from_str_radix(bin_digits, 2).map_err(|error| ParsingError::InvalidLiteral {
+                span,
+                kind: int_error_kind_to_literal_error_kind(
+                    error.kind(),
+                    LiteralErrorKind::FeltOverflow,
+                ),
+            })?;
+        if value > Felt::MODULUS {
+            return Err(ParsingError::InvalidLiteral {
+                span,
+                kind: LiteralErrorKind::FeltOverflow,
+            });
+        }
+        Ok(shrink_u64_bin(value))
+    } else {
+        Err(ParsingError::InvalidBinaryLiteral { span })
+    }
+}
+
+fn is_ascii_binary(c: char) -> bool {
+    matches!(c, '0'..='1')
+}
+
 #[inline]
-fn shrink_u64(n: u64) -> HexEncodedValue {
+fn shrink_u64_hex(n: u64) -> HexEncodedValue {
     if n <= (u8::MAX as u64) {
         HexEncodedValue::U8(n as u8)
     } else if n <= (u16::MAX as u64) {
@@ -622,6 +676,19 @@ fn shrink_u64(n: u64) -> HexEncodedValue {
     }
 }
 
+#[inline]
+fn shrink_u64_bin(n: u64) -> BinEncodedValue {
+    if n <= (u8::MAX as u64) {
+        BinEncodedValue::U8(n as u8)
+    } else if n <= (u16::MAX as u64) {
+        BinEncodedValue::U16(n as u16)
+    } else if n <= (u32::MAX as u64) {
+        BinEncodedValue::U32(n as u32)
+    } else {
+        BinEncodedValue::Felt(Felt::new(n))
+    }
+}
+
 #[inline]
 fn int_error_kind_to_literal_error_kind(
     kind: &IntErrorKind,