tasm_lib/io/
read_input.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use std::collections::HashMap;

use triton_vm::prelude::*;

use super::InputSource;
use crate::data_type::DataType;
use crate::empty_stack;
use crate::traits::basic_snippet::BasicSnippet;
use crate::traits::procedure::Procedure;
use crate::traits::procedure::ProcedureInitialState;

/// Move an element of type `DataType` from standard-in or secret-in's token stream to the stack
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct ReadInput {
    pub data_type: DataType,
    pub input_source: InputSource,
}

impl BasicSnippet for ReadInput {
    fn inputs(&self) -> Vec<(DataType, String)> {
        vec![]
    }

    fn outputs(&self) -> Vec<(DataType, String)> {
        vec![(self.data_type.clone(), "read_value".to_string())]
    }

    fn entrypoint(&self) -> String {
        format!(
            "tasmlib_io_read_{}___{}",
            self.input_source.label_friendly_name(),
            self.data_type.label_friendly_name()
        )
    }

    fn code(&self, _library: &mut crate::library::Library) -> Vec<LabelledInstruction> {
        let entrypoint = self.entrypoint();
        let read_an_element = self.data_type.read_value_from_input(self.input_source);
        triton_asm!(
            {entrypoint}:
                {&read_an_element}
                return
        )
    }
}

impl Procedure for ReadInput {
    fn rust_shadow(
        &self,
        stack: &mut Vec<BFieldElement>,
        _memory: &mut HashMap<BFieldElement, BFieldElement>,
        nondeterminism: &NonDeterminism,
        public_input: &[BFieldElement],
        _sponge: &mut Option<crate::VmHasher>,
    ) -> Vec<BFieldElement> {
        let input_source = match self.input_source {
            InputSource::StdIn => public_input,
            InputSource::SecretIn => &nondeterminism.individual_tokens,
        };
        for elem in input_source.iter().take(self.data_type.stack_size()) {
            stack.push(*elem);
        }

        // Output nothing
        vec![]
    }

    fn pseudorandom_initial_state(
        &self,
        _seed: [u8; 32],
        _bench_case: Option<crate::snippet_bencher::BenchmarkCase>,
    ) -> ProcedureInitialState {
        let input_stream: Vec<BFieldElement> = self.data_type.random_elements(1)[0].encode();

        let (std_in, secret_in) = match self.input_source {
            InputSource::StdIn => (input_stream, vec![]),
            InputSource::SecretIn => (vec![], input_stream),
        };

        ProcedureInitialState {
            stack: empty_stack(),
            nondeterminism: NonDeterminism::new(secret_in),
            public_input: std_in,
            sponge: None,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::traits::procedure::ShadowedProcedure;
    use crate::traits::rust_shadow::RustShadow;

    #[test]
    fn test() {
        for data_type in DataType::big_random_generatable_type_collection() {
            ShadowedProcedure::new(ReadInput {
                data_type: data_type.clone(),
                input_source: InputSource::StdIn,
            })
            .test();
            ShadowedProcedure::new(ReadInput {
                data_type,
                input_source: InputSource::SecretIn,
            })
            .test();
        }
    }
}

#[cfg(test)]
mod benches {
    use super::*;
    use crate::traits::procedure::ShadowedProcedure;
    use crate::traits::rust_shadow::RustShadow;

    #[test]
    fn bench_for_digest_reading() {
        ShadowedProcedure::new(ReadInput {
            data_type: DataType::Digest,
            input_source: InputSource::StdIn,
        })
        .bench();
    }
}