tasm_lib/
io.rs

1use std::fmt::Display;
2
3use triton_vm::prelude::*;
4
5pub mod read_input;
6pub mod write_to_stdout;
7
8#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
9pub enum InputSource {
10    StdIn,
11    SecretIn,
12}
13
14impl Display for InputSource {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        let str = match self {
17            InputSource::StdIn => "stdin",
18            InputSource::SecretIn => "secin",
19        };
20
21        write!(f, "{str}")
22    }
23}
24
25/// Return the code to write `n` words to output
26/// ```text
27/// BEFORE: _ [words; n]
28/// AFTER:  _
29/// ```
30pub fn write_words(n: usize) -> Vec<LabelledInstruction> {
31    let num_full_chunk_writes = n / 5;
32    let num_remaining_words = n % 5;
33    let mut instructions = vec![triton_instr!(write_io 5); num_full_chunk_writes];
34    if num_remaining_words > 0 {
35        instructions.extend(triton_asm!(write_io {
36            num_remaining_words
37        }));
38    }
39
40    instructions
41}
42
43impl InputSource {
44    /// The name of the instruction that reads from this input source
45    pub const fn instruction_name(&self) -> &str {
46        match self {
47            InputSource::StdIn => "read_io",
48            InputSource::SecretIn => "divine",
49        }
50    }
51
52    /// Return a string identifiying the input source and usable as assembly label
53    pub fn label_friendly_name(&self) -> &str {
54        match self {
55            InputSource::StdIn => "stdin",
56            InputSource::SecretIn => "secin",
57        }
58    }
59
60    /// Return the code used to read `n` words from the input source.
61    /// ```text
62    /// BEFORE: _
63    /// AFTER:  _ [read_words; n]
64    /// ```
65    pub fn read_words(&self, n: usize) -> Vec<LabelledInstruction> {
66        let input_instruction = self.instruction_name();
67
68        let num_full_chunk_reads = n / 5;
69        let num_remaining_words = n % 5;
70        let mut instructions =
71            vec![triton_asm!({input_instruction} 5); num_full_chunk_reads].concat();
72        if num_remaining_words > 0 {
73            instructions.extend(triton_asm!({input_instruction} {num_remaining_words}));
74        }
75
76        instructions
77    }
78}