circom-compat/src/witness/memory.rs
src/witness/memory.rsRaw
1//! Safe-ish interface for reading and writing specific types to the WASM runtime's memory
2use num_traits::ToPrimitive;
3use wasmer::{Memory, MemoryView};
4
5// TODO: Decide whether we want Ark here or if it should use a generic BigInt package
6use ark_bn254::FrParameters;
7use ark_ff::{BigInteger, BigInteger256, FpParameters, FromBytes, Zero};
8
9use num_bigint::{BigInt, BigUint};
10
11use color_eyre::Result;
12use std::str::FromStr;
13use std::{convert::TryFrom, ops::Deref};
14
15#[derive(Clone, Debug)]
16pub struct SafeMemory {
17 pub memory: Memory,
18 pub prime: BigInt,
19
20 short_max: BigInt,
21 short_min: BigInt,
22 r_inv: BigInt,
23 n32: usize,
24}
25
26impl Deref for SafeMemory {
27 type Target = Memory;
28
29 fn deref(&self) -> &Self::Target {
30 &self.memory
31 }
32}
33
34impl SafeMemory {
35 /// Creates a new SafeMemory
36 pub fn new(memory: Memory, n32: usize, prime: BigInt) -> Self {
37 // TODO: Figure out a better way to calculate these
38 let short_max = BigInt::from(0x8000_0000u64);
39 let short_min = BigInt::from_biguint(
40 num_bigint::Sign::NoSign,
41 BigUint::try_from(FrParameters::MODULUS).unwrap(),
42 ) - &short_max;
43 let r_inv = BigInt::from_str(
44 "9915499612839321149637521777990102151350674507940716049588462388200839649614",
45 )
46 .unwrap();
47
48 Self {
49 memory,
50 prime,
51
52 short_max,
53 short_min,
54 r_inv,
55 n32,
56 }
57 }
58
59 /// Gets an immutable view to the memory in 32 byte chunks
60 pub fn view(&self) -> MemoryView<u32> {
61 self.memory.view()
62 }
63
64 /// Returns the next free position in the memory
65 pub fn free_pos(&self) -> u32 {
66 self.view()[0].get()
67 }
68
69 /// Sets the next free position in the memory
70 pub fn set_free_pos(&mut self, ptr: u32) {
71 self.write_u32(0, ptr);
72 }
73
74 /// Allocates a U32 in memory
75 pub fn alloc_u32(&mut self) -> u32 {
76 let p = self.free_pos();
77 self.set_free_pos(p + 8);
78 p
79 }
80
81 /// Writes a u32 to the specified memory offset
82 pub fn write_u32(&mut self, ptr: usize, num: u32) {
83 let buf = unsafe { self.memory.data_unchecked_mut() };
84 buf[ptr..ptr + std::mem::size_of::<u32>()].copy_from_slice(&num.to_le_bytes());
85 }
86
87 /// Reads a u32 from the specified memory offset
88 pub fn read_u32(&self, ptr: usize) -> u32 {
89 let buf = unsafe { self.memory.data_unchecked() };
90
91 let mut bytes = [0; 4];
92 bytes.copy_from_slice(&buf[ptr..ptr + std::mem::size_of::<u32>()]);
93
94 u32::from_le_bytes(bytes)
95 }
96
97 /// Allocates `self.n32 * 4 + 8` bytes in the memory
98 pub fn alloc_fr(&mut self) -> u32 {
99 let p = self.free_pos();
100 self.set_free_pos(p + self.n32 as u32 * 4 + 8);
101 p
102 }
103
104 /// Writes a Field Element to memory at the specified offset, truncating
105 /// to smaller u32 types if needed and adjusting the sign via 2s complement
106 pub fn write_fr(&mut self, ptr: usize, fr: &BigInt) -> Result<()> {
107 if fr < &self.short_max && fr > &self.short_min {
108 if fr >= &BigInt::zero() {
109 self.write_short_positive(ptr, fr)?;
110 } else {
111 self.write_short_negative(ptr, fr)?;
112 }
113 } else {
114 self.write_long_normal(ptr, fr)?;
115 }
116
117 Ok(())
118 }
119
120 /// Reads a Field Element from the memory at the specified offset
121 pub fn read_fr(&self, ptr: usize) -> Result<BigInt> {
122 let view = self.memory.view::<u8>();
123
124 let res = if view[ptr + 4 + 3].get() & 0x80 != 0 {
125 let mut num = self.read_big(ptr + 8, self.n32)?;
126 if view[ptr + 4 + 3].get() & 0x40 != 0 {
127 num = (num * &self.r_inv) % &self.prime
128 }
129 num
130 } else if view[ptr + 3].get() & 0x40 != 0 {
131 let mut num = self.read_u32(ptr).into();
132 // handle small negative
133 num -= BigInt::from(0x100000000i64);
134 num
135 } else {
136 self.read_u32(ptr).into()
137 };
138
139 Ok(res)
140 }
141
142 fn write_short_positive(&mut self, ptr: usize, fr: &BigInt) -> Result<()> {
143 let num = fr.to_i32().expect("not a short positive");
144 self.write_u32(ptr, num as u32);
145 self.write_u32(ptr + 4, 0);
146 Ok(())
147 }
148
149 fn write_short_negative(&mut self, ptr: usize, fr: &BigInt) -> Result<()> {
150 // 2s complement
151 let num = fr - &self.short_min;
152 let num = num - &self.short_max;
153 let num = num + BigInt::from(0x0001_0000_0000i64);
154
155 let num = num
156 .to_u32()
157 .expect("could not cast as u32 (should never happen)");
158
159 self.write_u32(ptr, num);
160 self.write_u32(ptr + 4, 0);
161 Ok(())
162 }
163
164 fn write_long_normal(&mut self, ptr: usize, fr: &BigInt) -> Result<()> {
165 self.write_u32(ptr, 0);
166 self.write_u32(ptr + 4, i32::MIN as u32); // 0x80000000
167 self.write_big(ptr + 8, fr)?;
168 Ok(())
169 }
170
171 fn write_big(&self, ptr: usize, num: &BigInt) -> Result<()> {
172 let buf = unsafe { self.memory.data_unchecked_mut() };
173
174 // TODO: How do we handle negative bignums?
175 let (_, num) = num.clone().into_parts();
176 let num = BigInteger256::try_from(num).unwrap();
177
178 let bytes = num.to_bytes_le();
179 let len = bytes.len();
180 buf[ptr..ptr + len].copy_from_slice(&bytes);
181
182 Ok(())
183 }
184
185 /// Reads `num_bytes * 32` from the specified memory offset in a Big Integer
186 pub fn read_big(&self, ptr: usize, num_bytes: usize) -> Result<BigInt> {
187 let buf = unsafe { self.memory.data_unchecked() };
188 let buf = &buf[ptr..ptr + num_bytes * 32];
189
190 // TODO: Is there a better way to read big integers?
191 let big = BigInteger256::read(buf).unwrap();
192 let big = BigUint::try_from(big).unwrap();
193 Ok(big.into())
194 }
195}
196
197// TODO: Figure out how to read / write numbers > u32
198// circom-witness-calculator: Wasm + Memory -> expose BigInts so that they can be consumed by any proof system
199// ark-circom:
200// 1. can read zkey
201// 2. can generate witness from inputs
202// 3. can generate proofs
203// 4. can serialize proofs in the desired format
204#[cfg(test)]
205mod tests {
206 use super::*;
207 use num_traits::ToPrimitive;
208 use std::str::FromStr;
209 use wasmer::{MemoryType, Store};
210
211 fn new() -> SafeMemory {
212 SafeMemory::new(
213 Memory::new(&Store::default(), MemoryType::new(1, None, false)).unwrap(),
214 2,
215 BigInt::from_str(
216 "21888242871839275222246405745257275088548364400416034343698204186575808495617",
217 )
218 .unwrap(),
219 )
220 }
221
222 #[test]
223 fn i32_bounds() {
224 let mem = new();
225 let i32_max = i32::MAX as i64 + 1;
226 assert_eq!(mem.short_min.to_i64().unwrap(), -i32_max);
227 assert_eq!(mem.short_max.to_i64().unwrap(), i32_max);
228 }
229
230 #[test]
231 fn read_write_32() {
232 let mut mem = new();
233 let num = u32::MAX;
234
235 let inp = mem.read_u32(0);
236 assert_eq!(inp, 0);
237
238 mem.write_u32(0, num);
239 let inp = mem.read_u32(0);
240 assert_eq!(inp, num);
241 }
242
243 #[test]
244 fn read_write_fr_small_positive() {
245 read_write_fr(BigInt::from(1_000_000));
246 }
247
248 #[test]
249 fn read_write_fr_small_negative() {
250 read_write_fr(BigInt::from(-1_000_000));
251 }
252
253 #[test]
254 fn read_write_fr_big_positive() {
255 read_write_fr(BigInt::from(500000000000i64));
256 }
257
258 // TODO: How should this be handled?
259 #[test]
260 #[ignore]
261 fn read_write_fr_big_negative() {
262 read_write_fr(BigInt::from_str("-500000000000").unwrap())
263 }
264
265 fn read_write_fr(num: BigInt) {
266 let mut mem = new();
267 mem.write_fr(0, &num).unwrap();
268 let res = mem.read_fr(0).unwrap();
269 assert_eq!(res, num);
270 }
271}