feat!: make instruction write_io take an argument in range 1..=5

BREAKING CHANGE: instruction `write_io` takes an argument

triton-vm/src/error.rs

@@ -20,6 +20,7 @@ pub enum InstructionError {
     IllegalPop(usize),
     IllegalDivine(usize),
     IllegalReadIo(usize),
+    IllegalWriteIo(usize),
     SwapST0,
     InverseOfZero,
     DivisionByZero,
@@ -50,6 +51,9 @@ impl Display for InstructionError {
             IllegalPop(st) => write!(f, "must pop at least 1, at most 5 elements, not {st}"),
             IllegalDivine(st) => write!(f, "must divine at least 1, at most 5 elements, not {st}"),
             IllegalReadIo(st) => write!(f, "must read_io at least 1, at most 5 elements, not {st}"),
+            IllegalWriteIo(st) => {
+                write!(f, "must write_io at least 1, at most 5 elements, not {st}")
+            }
             SwapST0 => write!(f, "Cannot swap stack element 0 with itself"),
             InverseOfZero => write!(f, "0 does not have a multiplicative inverse"),
             DivisionByZero => write!(f, "Division by 0 is impossible"),

triton-vm/src/example_programs.rs

@@ -29,7 +29,7 @@ fn fibonacci_sequence() -> Program {
 
         // pop zero, write result
         pop 1
-        write_io
+        write_io 1
         halt
 
         // before: ⊥ 0 1 i
@@ -74,7 +74,7 @@ fn greatest_common_divisor() -> Program {
 
         terminate:
             // _ d n where d == 0
-            write_io // _ d
+            write_io 1      // _ d
             halt
     )
 }

triton-vm/src/instruction.rs

@@ -149,7 +149,7 @@ pub enum AnInstruction<Dest: PartialEq + Default> {
 
     // Read/write
     ReadIo(OpStackElement),
-    WriteIo,
+    WriteIo(OpStackElement),
 }
 
 impl<Dest: PartialEq + Default> AnInstruction<Dest> {
@@ -193,7 +193,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             XInvert => 80,
             XbMul => 82,
             ReadIo(_) => 41,
-            WriteIo => 90,
+            WriteIo(_) => 11,
         }
     }
 
@@ -236,7 +236,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             XInvert => "xinvert",
             XbMul => "xbmul",
             ReadIo(_) => "read_io",
-            WriteIo => "write_io",
+            WriteIo(_) => "write_io",
         }
     }
 
@@ -250,7 +250,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             Divine(_) => 2,
             Dup(_) | Swap(_) => 2,
             Call(_) => 2,
-            ReadIo(_) => 2,
+            ReadIo(_) | WriteIo(_) => 2,
             _ => 1,
         }
     }
@@ -306,7 +306,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             XInvert => XInvert,
             XbMul => XbMul,
             ReadIo(x) => ReadIo(*x),
-            WriteIo => WriteIo,
+            WriteIo(x) => WriteIo(*x),
         }
     }
 
@@ -361,7 +361,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             XInvert => 0,
             XbMul => -1,
             ReadIo(n) => n.index() as i32,
-            WriteIo => -1,
+            WriteIo(n) => -(n.index() as i32),
         }
     }
 
@@ -379,6 +379,7 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
             Divine(st) => *st < ST1 || ST5 < *st,
             Swap(ST0) => true,
             ReadIo(st) => *st < ST1 || ST5 < *st,
+            WriteIo(st) => *st < ST1 || ST5 < *st,
             _ => false,
         }
     }
@@ -389,8 +390,9 @@ impl<Dest: Display + PartialEq + Default> Display for AnInstruction<Dest> {
         write!(f, "{}", self.name())?;
         match self {
             Push(arg) => write!(f, " {arg}"),
-            Pop(arg) | Divine(arg) | Dup(arg) | Swap(arg) | ReadIo(arg) => write!(f, " {arg}"),
+            Pop(arg) | Divine(arg) | Dup(arg) | Swap(arg) => write!(f, " {arg}"),
             Call(arg) => write!(f, " {arg}"),
+            ReadIo(arg) | WriteIo(arg) => write!(f, " {arg}"),
             _ => Ok(()),
         }
     }
@@ -401,7 +403,8 @@ impl Instruction {
     pub fn arg(&self) -> Option<BFieldElement> {
         match self {
             Push(arg) | Call(arg) => Some(*arg),
-            Pop(arg) | Divine(arg) | Dup(arg) | Swap(arg) | ReadIo(arg) => Some(arg.into()),
+            Pop(arg) | Divine(arg) | Dup(arg) | Swap(arg) => Some(arg.into()),
+            ReadIo(arg) | WriteIo(arg) => Some(arg.into()),
             _ => None,
         }
     }
@@ -432,6 +435,7 @@ impl Instruction {
             Dup(_) => Some(Dup(stack_element)),
             Swap(_) => Some(Swap(stack_element)),
             ReadIo(_) => Some(ReadIo(stack_element)),
+            WriteIo(_) => Some(WriteIo(stack_element)),
             _ => None,
         };
         if new_instruction?.has_illegal_argument() {
@@ -518,7 +522,7 @@ const fn all_instructions_without_args() -> [AnInstruction<BFieldElement>; Instr
         XInvert,
         XbMul,
         ReadIo(ST0),
-        WriteIo,
+        WriteIo(ST0),
     ]
 }
 
@@ -633,6 +637,7 @@ mod tests {
                 Divine(ST0) => Divine(ST1),
                 Swap(ST0) => Swap(ST1),
                 ReadIo(ST0) => ReadIo(ST1),
+                WriteIo(ST0) => WriteIo(ST1),
                 _ => self,
             }
         }
@@ -856,7 +861,7 @@ mod tests {
 
     #[test]
     fn stringify_some_instructions() {
-        let instructions = triton_asm!(push 3 invert push 2 mul push 1 add write_io halt);
+        let instructions = triton_asm!(push 3 invert push 2 mul push 1 add write_io 1 halt);
         let code = stringify_instructions(&instructions);
         println!("{code}");
     }

triton-vm/src/lib.rs

@@ -34,7 +34,7 @@
 //!     read_io 1           // n
 //!     push 1              // n 1
 //!     call factorial      // 0 n!
-//!     write_io            // 0
+//!     write_io 1          // 0
 //!     halt
 //!
 //!     factorial:          // n acc
@@ -170,15 +170,15 @@ pub mod vm;
 /// let program = triton_program!(
 ///     read_io 1 push 5 mul
 ///     call check_eq_15
-///     push 1 write_io
+///     push 17 write_io 1
 ///     halt
 ///     // assert that the top of the stack is 15
 ///     check_eq_15:
 ///         push 15 eq assert
 ///         return
 /// );
 /// let output = program.run(vec![3].into(), [].into()).unwrap();
-/// assert_eq!(1, output[0].value());
+/// assert_eq!(17, output[0].value());
 /// ```
 ///
 /// Any type with an appropriate [`Display`](std::fmt::Display) implementation can be
@@ -342,6 +342,7 @@ macro_rules! triton_asm {
     [swap $arg:literal; $num:expr] => { vec![ $crate::triton_instr!(swap $arg); $num ] };
     [call $arg:ident; $num:expr] => { vec![ $crate::triton_instr!(call $arg); $num ] };
     [read_io $arg:literal; $num:expr] => { vec![ $crate::triton_instr!(read_io $arg); $num ] };
+    [write_io $arg:literal; $num:expr] => { vec![ $crate::triton_instr!(write_io $arg); $num ] };
     [$instr:ident; $num:expr] => { vec![ $crate::triton_instr!($instr); $num ] };
 
     // entry point
@@ -407,6 +408,12 @@ macro_rules! triton_instr {
         let instruction = $crate::instruction::AnInstruction::<String>::ReadIo(argument);
         $crate::instruction::LabelledInstruction::Instruction(instruction)
     }};
+    (write_io $arg:literal) => {{
+        assert!(1_u32 <= $arg && $arg <= 5, "`write_io {}` is illegal", $arg);
+        let argument: $crate::op_stack::OpStackElement = u32::try_into($arg).unwrap();
+        let instruction = $crate::instruction::AnInstruction::<String>::WriteIo(argument);
+        $crate::instruction::LabelledInstruction::Instruction(instruction)
+    }};
     ($instr:ident) => {{
         let (_, instructions) = $crate::parser::tokenize(stringify!($instr)).unwrap();
         instructions[0].to_labelled_instruction()
@@ -606,7 +613,7 @@ mod tests {
             push 51 read_mem
             push 42 read_mem
             swap 1 swap 2 mul
-            write_io halt
+            write_io 1 halt
         );
 
         let initial_ram = [(42, 17), (51, 13)].into();
@@ -740,15 +747,16 @@ mod tests {
 
     #[test]
     fn nested_triton_asm_interpolation() {
-        let triple_write = triton_asm![write_io; 3];
+        let double_write = triton_asm![write_io 1; 2];
+        let quadruple_write = triton_asm!({&double_write} write_io 2);
         let snippet_0 = triton_asm!(push 7 nop call my_label);
-        let snippet_1 = triton_asm!(pop 2 halt my_label: push 8 {&triple_write});
+        let snippet_1 = triton_asm!(pop 2 halt my_label: push 8 push 9 {&quadruple_write});
         let source_code = triton_asm!(push 6 {&snippet_0} {&snippet_1} halt);
 
         let program = triton_program!({ &source_code });
         let public_output = program.run([].into(), [].into()).unwrap();
 
-        let expected_output = [8, 7, 6].map(BFieldElement::new).to_vec();
+        let expected_output = [9, 8, 7, 6].map(BFieldElement::new).to_vec();
         assert_eq!(expected_output, public_output);
     }
 

triton-vm/src/parser.rs

@@ -276,7 +276,7 @@ fn an_instruction(s: &str) -> ParseResult<AnInstruction<String>> {
 
     // Read/write
     let read_io = read_io_instruction();
-    let write_io = instruction("write_io", WriteIo);
+    let write_io = write_io_instruction();
 
     let read_write = alt((read_io, write_io));
 
@@ -423,6 +423,21 @@ fn read_io_instruction() -> impl Fn(&str) -> ParseResult<AnInstruction<String>>
     }
 }
 
+fn write_io_instruction() -> impl Fn(&str) -> ParseResult<AnInstruction<String>> {
+    move |s: &str| {
+        let (s, _) = token1("write_io")(s)?; // require space after instruction name
+        let (s, stack_register) = stack_register(s)?;
+        let (s, _) = comment_or_whitespace1(s)?;
+
+        let instruction = WriteIo(stack_register);
+        if instruction.has_illegal_argument() {
+            return cut(context("illegal argument for instruction `write_io`", fail))(s);
+        }
+
+        Ok((s, instruction))
+    }
+}
+
 fn field_element(s_orig: &str) -> ParseResult<BFieldElement> {
     let (s, negative) = opt(token0("-"))(s_orig)?;
     let (s, n) = digit1(s)?;
@@ -655,7 +670,7 @@ pub(crate) mod tests {
         let mut rng = thread_rng();
 
         let difficult_instructions = vec![
-            "pop", "push", "divine", "dup", "swap", "skiz", "call", "read_io",
+            "pop", "push", "divine", "dup", "swap", "skiz", "call", "read_io", "write_io",
         ];
         let simple_instructions = ALL_INSTRUCTION_NAMES
             .into_iter()
@@ -664,7 +679,7 @@ pub(crate) mod tests {
 
         let generators = [vec!["simple"], difficult_instructions].concat();
         // Test difficult instructions more frequently.
-        let weights = vec![simple_instructions.len(), 2, 2, 2, 6, 6, 2, 10, 2];
+        let weights = vec![simple_instructions.len(), 2, 2, 2, 6, 6, 2, 10, 2, 2];
 
         assert_eq!(
             generators.len(),
@@ -722,6 +737,11 @@ pub(crate) mod tests {
                 vec!["read_io".to_string(), format!("{arg}")]
             }
 
+            "write_io" => {
+                let arg: usize = rng.gen_range(1..=5);
+                vec!["write_io".to_string(), format!("{arg}")]
+            }
+
             unknown => panic!("Unknown generator, {unknown}"),
         }
     }
@@ -1020,8 +1040,8 @@ pub(crate) mod tests {
 
     #[test]
     fn triton_asm_macro() {
-        let instructions = triton_asm!(write_io push 17 call which_label lt swap 3);
-        assert_eq!(Instruction(WriteIo), instructions[0]);
+        let instructions = triton_asm!(write_io 3 push 17 call which_label lt swap 3);
+        assert_eq!(Instruction(WriteIo(ST3)), instructions[0]);
         assert_eq!(Instruction(Push(17_u64.into())), instructions[1]);
         assert_eq!(
             Instruction(Call("which_label".to_string())),

triton-vm/src/stark.rs

@@ -1416,7 +1416,7 @@ pub(crate) mod tests {
     #[test]
     fn check_io_terminals() {
         let read_nop_program = triton_program!(
-            read_io 3 nop nop write_io push 17 write_io halt
+            read_io 3 nop nop write_io 2 push 17 write_io 1 halt
         );
         let mut program_and_input = ProgramAndInput::without_input(read_nop_program);
         program_and_input.public_input = vec![3, 5, 7];