tasm_lib/io/
read_input.rs

1use std::collections::HashMap;
2
3use triton_vm::prelude::*;
4
5use super::InputSource;
6use crate::empty_stack;
7use crate::prelude::*;
8use crate::traits::procedure::Procedure;
9use crate::traits::procedure::ProcedureInitialState;
10
11/// Move an element of type `DataType` from standard-in or secret-in's token stream to the stack
12#[derive(Debug, Clone, Eq, PartialEq, Hash)]
13pub struct ReadInput {
14    pub data_type: DataType,
15    pub input_source: InputSource,
16}
17
18impl BasicSnippet for ReadInput {
19    fn inputs(&self) -> Vec<(DataType, String)> {
20        vec![]
21    }
22
23    fn outputs(&self) -> Vec<(DataType, String)> {
24        vec![(self.data_type.clone(), "read_value".to_string())]
25    }
26
27    fn entrypoint(&self) -> String {
28        format!(
29            "tasmlib_io_read_{}___{}",
30            self.input_source.label_friendly_name(),
31            self.data_type.label_friendly_name()
32        )
33    }
34
35    fn code(&self, _library: &mut crate::library::Library) -> Vec<LabelledInstruction> {
36        let entrypoint = self.entrypoint();
37        let read_an_element = self.data_type.read_value_from_input(self.input_source);
38        triton_asm!(
39            {entrypoint}:
40                {&read_an_element}
41                return
42        )
43    }
44}
45
46impl Procedure for ReadInput {
47    fn rust_shadow(
48        &self,
49        stack: &mut Vec<BFieldElement>,
50        _memory: &mut HashMap<BFieldElement, BFieldElement>,
51        nondeterminism: &NonDeterminism,
52        public_input: &[BFieldElement],
53        _sponge: &mut Option<crate::prelude::Tip5>,
54    ) -> Vec<BFieldElement> {
55        let input_source = match self.input_source {
56            InputSource::StdIn => public_input,
57            InputSource::SecretIn => &nondeterminism.individual_tokens,
58        };
59        for elem in input_source.iter().take(self.data_type.stack_size()) {
60            stack.push(*elem);
61        }
62
63        // Output nothing
64        vec![]
65    }
66
67    fn pseudorandom_initial_state(
68        &self,
69        _seed: [u8; 32],
70        _bench_case: Option<crate::snippet_bencher::BenchmarkCase>,
71    ) -> ProcedureInitialState {
72        let input_stream: Vec<BFieldElement> = self.data_type.random_elements(1)[0].encode();
73
74        let (std_in, secret_in) = match self.input_source {
75            InputSource::StdIn => (input_stream, vec![]),
76            InputSource::SecretIn => (vec![], input_stream),
77        };
78
79        ProcedureInitialState {
80            stack: empty_stack(),
81            nondeterminism: NonDeterminism::new(secret_in),
82            public_input: std_in,
83            sponge: None,
84        }
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::test_prelude::*;
92
93    #[test]
94    fn test() {
95        for data_type in DataType::big_random_generatable_type_collection() {
96            ShadowedProcedure::new(ReadInput {
97                data_type: data_type.clone(),
98                input_source: InputSource::StdIn,
99            })
100            .test();
101            ShadowedProcedure::new(ReadInput {
102                data_type,
103                input_source: InputSource::SecretIn,
104            })
105            .test();
106        }
107    }
108}
109
110#[cfg(test)]
111mod benches {
112    use super::*;
113    use crate::test_prelude::*;
114
115    #[test]
116    fn benchmark() {
117        ShadowedProcedure::new(ReadInput {
118            data_type: DataType::Digest,
119            input_source: InputSource::StdIn,
120        })
121        .bench();
122    }
123}