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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

//! The flat vector is a vector of fixed-size elements that we want to expose directly to js-of-ocaml
//! (so that we can access a `Vec<Field>` cheaply,
//! by just passing a pointer to a continuous memory region instead of copying.
//! The wasmvector is a normal heap-allocated vector,
//! where we leave it on the rust heap and just keep a pointer around.
//! We use flat for fields, normal for gates etc.
//!
//! Accessing Rust vector values is not the same as accessing an array.
//! Each indexing (e.g. `some_vec[3]`) is costly as it is implemented as a function call.
//! Knowing that, plus the fact that field elements are implemented as `[u32; 8]`, we know that we incur the cost of following several pointers.
//! To decrease that cost, we flatten such arrays, going from something like
//!
//! ```ignore
//! [[a0, a1, ..., a7], [b0, b1, ..., b7], ...]
//! ```
//!
//! to a flattened vector like:
//!
//! ```ignore
//! [a0, a1, ..., a7, b0, b1, ..., b7, ...]
//! ```

use wasm_bindgen::convert::{FromWasmAbi, IntoWasmAbi, OptionFromWasmAbi, OptionIntoWasmAbi};

use core::{convert::From, ops::Deref};

#[derive(Clone, Debug)]
pub struct WasmFlatVector<T>(Vec<T>);

pub trait FlatVectorElem {
    const FLATTENED_SIZE: usize;
    fn flatten(self) -> Vec<u8>;
    fn unflatten(flat: Vec<u8>) -> Self;
}

impl<T> Deref for WasmFlatVector<T> {
    type Target = Vec<T>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> From<Vec<T>> for WasmFlatVector<T> {
    fn from(x: Vec<T>) -> Self {
        WasmFlatVector(x)
    }
}

impl<T> From<WasmFlatVector<T>> for Vec<T> {
    fn from(x: WasmFlatVector<T>) -> Self {
        x.0
    }
}

impl<'a, T> From<&'a WasmFlatVector<T>> for &'a Vec<T> {
    fn from(x: &'a WasmFlatVector<T>) -> Self {
        &x.0
    }
}

impl<T> core::iter::IntoIterator for WasmFlatVector<T> {
    type Item = <Vec<T> as core::iter::IntoIterator>::Item;
    type IntoIter = <Vec<T> as core::iter::IntoIterator>::IntoIter;
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<'a, T> core::iter::IntoIterator for &'a WasmFlatVector<T> {
    type Item = <&'a Vec<T> as core::iter::IntoIterator>::Item;
    type IntoIter = <&'a Vec<T> as core::iter::IntoIterator>::IntoIter;
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<T> core::iter::FromIterator<T> for WasmFlatVector<T> {
    fn from_iter<I>(iter: I) -> WasmFlatVector<T>
    where
        I: IntoIterator<Item = T>,
    {
        WasmFlatVector(core::iter::FromIterator::from_iter(iter))
    }
}

impl<T> core::default::Default for WasmFlatVector<T> {
    fn default() -> Self {
        WasmFlatVector(core::default::Default::default())
    }
}

impl<T> core::iter::Extend<T> for WasmFlatVector<T> {
    fn extend<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = T>,
    {
        self.0.extend(iter)
    }
}

impl<T> wasm_bindgen::describe::WasmDescribe for WasmFlatVector<T> {
    fn describe() {
        <Vec<u8> as wasm_bindgen::describe::WasmDescribe>::describe()
    }
}

impl<T: FlatVectorElem> FromWasmAbi for WasmFlatVector<T> {
    type Abi = <Vec<u8> as FromWasmAbi>::Abi;
    unsafe fn from_abi(js: Self::Abi) -> Self {
        let data: Vec<u8> = FromWasmAbi::from_abi(js);
        let mut res: Vec<T> = Vec::with_capacity(data.len() / T::FLATTENED_SIZE);

        let mut buf = Vec::with_capacity(T::FLATTENED_SIZE);
        for x in data.into_iter() {
            assert!(buf.len() < T::FLATTENED_SIZE);
            buf.push(x);
            if buf.len() >= T::FLATTENED_SIZE {
                res.push(T::unflatten(buf));
                buf = Vec::with_capacity(T::FLATTENED_SIZE);
            }
        }
        assert_eq!(buf.len(), 0);
        WasmFlatVector(res)
    }
}

impl<T: FlatVectorElem> OptionFromWasmAbi for WasmFlatVector<T> {
    fn is_none(x: &Self::Abi) -> bool {
        <Vec<u8> as OptionFromWasmAbi>::is_none(x)
    }
}

impl<T: FlatVectorElem> IntoWasmAbi for WasmFlatVector<T> {
    type Abi = <Vec<u8> as FromWasmAbi>::Abi;
    fn into_abi(self) -> Self::Abi {
        let mut data: Vec<u8> = Vec::with_capacity(self.0.len() * T::FLATTENED_SIZE);
        for x in self.0.into_iter() {
            data.extend(x.flatten().into_iter());
        }
        IntoWasmAbi::into_abi(data)
    }
}

impl<T: FlatVectorElem> OptionIntoWasmAbi for WasmFlatVector<T> {
    fn none() -> Self::Abi {
        <Vec<u8> as OptionIntoWasmAbi>::none()
    }
}