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}