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