accept instructions with argument in macro triton_instr!()

TritonVM · Jul 12, 2023 · 15d9b68 · 15d9b68
1 parent 6ecb3f5
commit 15d9b68
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/triton-vm/src/lib.rs b/triton-vm/src/lib.rs
@@ -166,12 +166,32 @@ macro_rules! triton_asm {
     }};
 }
 
-/// Compile a single [Triton assembly][tasm] instruction. Output a
-/// [`LabelledInstruction`].
+/// Compile a single [Triton assembly][tasm] instruction into a [`LabelledInstruction`].
 ///
 /// [tasm]: https://triton-vm.org/spec/instructions.html
 #[macro_export]
 macro_rules! triton_instr {
+    (push $arg:literal) => {{
+        let argument = $crate::BFieldElement::new($arg);
+        let instruction = $crate::instruction::AnInstruction::<String>::Push(argument);
+        $crate::instruction::LabelledInstruction::Instruction(instruction)
+    }};
+    (dup $arg:literal) => {{
+        let argument: $crate::op_stack::OpStackElement = u32::try_into($arg).unwrap();
+        let instruction = $crate::instruction::AnInstruction::<String>::Dup(argument);
+        $crate::instruction::LabelledInstruction::Instruction(instruction)
+    }};
+    (swap $arg:literal) => {{
+        assert_ne!(0_u32, $arg, "`swap 0` is illegal.");
+        let argument: $crate::op_stack::OpStackElement = u32::try_into($arg).unwrap();
+        let instruction = $crate::instruction::AnInstruction::<String>::Swap(argument);
+        $crate::instruction::LabelledInstruction::Instruction(instruction)
+    }};
+    (call $arg:ident) => {{
+        let argument = stringify!($arg).to_string();
+        let instruction = $crate::instruction::AnInstruction::<String>::Call(argument);
+        $crate::instruction::LabelledInstruction::Instruction(instruction)
+    }};
     ($instr:ident) => {{
         let (_, instructions) = $crate::parser::tokenize(stringify!($instr)).unwrap();
         instructions[0].to_labelled_instruction()

diff --git a/triton-vm/src/parser.rs b/triton-vm/src/parser.rs
@@ -1017,7 +1017,26 @@ pub mod parser_tests {
     }
 
     #[test]
-    fn triton_instruction_macro() {
-        assert_eq!(LabelledInstruction::Instruction(Halt), triton_instr!(halt));
+    fn test_triton_instruction_macro_on_simple_instructions() {
+        assert_eq!(Instruction(Halt), triton_instr!(halt));
+        assert_eq!(Instruction(Add), triton_instr!(add));
+        assert_eq!(Instruction(Pop), triton_instr!(pop));
+    }
+
+    #[test]
+    #[should_panic(expected = "not_an_instruction")]
+    fn negative_test_of_triton_instruction_macro() {
+        triton_instr!(not_an_instruction);
+    }
+
+    #[test]
+    fn test_triton_instruction_macro_on_instructions_with_argument() {
+        assert_eq!(Instruction(Push(7_u64.into())), triton_instr!(push 7));
+        assert_eq!(Instruction(Dup(ST3)), triton_instr!(dup 3));
+        assert_eq!(Instruction(Swap(ST5)), triton_instr!(swap 5));
+        assert_eq!(
+            Instruction(Call("my_label".to_string())),
+            triton_instr!(call my_label)
+        );
     }
 }