openvm_sdk/
stdin.rs

1use std::{
2    collections::{HashMap, VecDeque},
3    sync::Arc,
4};
5
6use openvm_circuit::arch::Streams;
7use openvm_stark_backend::p3_field::Field;
8use serde::{Deserialize, Serialize};
9
10#[derive(Clone, Default, Serialize, Deserialize)]
11pub struct StdIn<F = crate::F> {
12    pub buffer: VecDeque<Vec<F>>,
13    pub kv_store: HashMap<Vec<u8>, Vec<u8>>,
14}
15
16impl<F: Field> StdIn<F> {
17    pub fn from_bytes(data: &[u8]) -> Self {
18        let mut ret = Self::default();
19        ret.write_bytes(data);
20        ret
21    }
22
23    pub fn read(&mut self) -> Option<Vec<F>> {
24        self.buffer.pop_front()
25    }
26
27    pub fn write<T: Serialize>(&mut self, data: &T) {
28        let words = openvm::serde::to_vec(data).unwrap();
29        let bytes: Vec<u8> = words.into_iter().flat_map(|w| w.to_le_bytes()).collect();
30        self.write_bytes(&bytes);
31    }
32
33    pub fn write_bytes(&mut self, data: &[u8]) {
34        let field_data = data.iter().map(|b| F::from_canonical_u8(*b)).collect();
35        self.buffer.push_back(field_data);
36    }
37
38    pub fn write_field(&mut self, data: &[F]) {
39        self.buffer.push_back(data.to_vec());
40    }
41    pub fn add_key_value(&mut self, key: Vec<u8>, value: Vec<u8>) {
42        self.kv_store.insert(key, value);
43    }
44}
45
46impl<F: Field> From<StdIn<F>> for Streams<F> {
47    fn from(mut std_in: StdIn<F>) -> Self {
48        let mut data = Vec::<Vec<F>>::new();
49        while let Some(input) = std_in.read() {
50            data.push(input);
51        }
52        let mut ret = Streams::new(data);
53        ret.kv_store = Arc::new(std_in.kv_store);
54        ret
55    }
56}
57
58impl<F: Field> From<Vec<Vec<F>>> for StdIn<F> {
59    fn from(inputs: Vec<Vec<F>>) -> Self {
60        let mut ret = StdIn::<F>::default();
61        for input in inputs {
62            ret.write_field(&input);
63        }
64        ret
65    }
66}