Skip to main content

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