1use std::io;
2
3use itertools::Itertools;
4
5use super::{
6 inner::{FqSlice, FqSliceMut, FqVector},
7 iter::{FqVectorIterator, FqVectorNonZeroIterator},
8};
9use crate::{
10 field::{Field, element::FieldElement},
11 limb::Limb,
12 prime::{Prime, ValidPrime},
13};
14
15impl<F: Field> FqVector<F> {
16 pub fn new(fq: F, len: usize) -> Self {
17 let number_of_limbs = fq.number(len);
18
19 Self::from_raw_parts(fq, len, vec![0; number_of_limbs])
20 }
21
22 pub fn new_with_capacity(fq: F, len: usize, capacity: usize) -> Self {
23 let mut limbs = Vec::with_capacity(fq.number(capacity));
24 limbs.resize(fq.number(len), 0);
25
26 Self::from_raw_parts(fq, len, limbs)
27 }
28
29 pub fn from_slice(fq: F, slice: &[FieldElement<F>]) -> Self {
30 assert!(slice.iter().all(|x| x.field() == fq));
31 let len = slice.len();
32 let mut v = Self::new(fq, len);
33 v.copy_from_slice(slice);
34 v
35 }
36
37 pub fn from_bytes(fq: F, len: usize, data: &mut impl io::Read) -> io::Result<Self> {
38 let mut v = Self::new(fq, len);
39 v.update_from_bytes(data)?;
40 Ok(v)
41 }
42
43 pub fn update_from_bytes(&mut self, data: &mut impl io::Read) -> io::Result<()> {
44 crate::limb::from_bytes(self.limbs_mut(), data)
45 }
46
47 pub fn to_bytes(&self, buffer: &mut impl io::Write) -> io::Result<()> {
48 let num_limbs = self.fq().number(self.len());
51 crate::limb::to_bytes(&self.limbs()[..num_limbs], buffer)
52 }
53
54 pub const fn is_empty(&self) -> bool {
55 self.len() == 0
56 }
57
58 pub fn prime(&self) -> ValidPrime {
59 self.fq().characteristic().to_dyn()
60 }
61
62 #[must_use]
63 pub fn slice(&self, start: usize, end: usize) -> FqSlice<'_, F> {
64 assert!(start <= end && end <= self.len());
65 FqSlice::new(self.fq(), self.limbs(), start, end)
66 }
67
68 #[must_use]
69 pub fn slice_mut(&mut self, start: usize, end: usize) -> FqSliceMut<'_, F> {
70 assert!(start <= end && end <= self.len());
71 FqSliceMut::new(self.fq(), self.limbs_mut(), start, end)
72 }
73
74 #[inline]
75 #[must_use]
76 pub fn as_slice(&self) -> FqSlice<'_, F> {
77 self.into()
78 }
79
80 #[inline]
81 #[must_use]
82 pub fn as_slice_mut(&mut self) -> FqSliceMut<'_, F> {
83 self.into()
84 }
85
86 pub fn add_basis_element(&mut self, index: usize, value: FieldElement<F>) {
87 assert_eq!(self.fq(), value.field());
88 self.as_slice_mut().add_basis_element(index, value);
89 }
90
91 pub fn entry(&self, index: usize) -> FieldElement<F> {
92 self.as_slice().entry(index)
93 }
94
95 pub fn set_entry(&mut self, index: usize, value: FieldElement<F>) {
96 assert_eq!(self.fq(), value.field());
97 self.as_slice_mut().set_entry(index, value);
98 }
99
100 pub fn iter(&self) -> FqVectorIterator<'_, F> {
101 self.as_slice().iter()
102 }
103
104 pub fn iter_nonzero(&self) -> FqVectorNonZeroIterator<'_, F> {
105 self.as_slice().iter_nonzero()
106 }
107
108 pub fn set_to_zero(&mut self) {
109 for limb in self.limbs_mut() {
111 *limb = 0;
112 }
113 }
114
115 pub fn scale(&mut self, c: FieldElement<F>) {
116 assert_eq!(self.fq(), c.field());
117 let fq = self.fq();
118
119 if c == fq.zero() {
120 self.set_to_zero();
121 }
122 if fq.q() != 2 {
123 for limb in self.limbs_mut() {
124 *limb = fq.reduce(fq.fma_limb(0, *limb, c.clone()));
125 }
126 }
127 }
128
129 pub fn add_offset(&mut self, other: &Self, c: FieldElement<F>, offset: usize) {
132 assert_eq!(self.fq(), c.field());
133 assert_eq!(self.fq(), other.fq());
134 assert_eq!(self.len(), other.len());
135
136 let fq = self.fq();
137 let min_limb = offset / fq.entries_per_limb();
138
139 if fq.q() == 2 {
140 if c != fq.zero() {
141 crate::simd::add_simd(self.limbs_mut(), other.limbs(), min_limb);
142 }
143 } else {
144 for (left, right) in self
145 .limbs_mut()
146 .iter_mut()
147 .zip_eq(other.limbs())
148 .skip(min_limb)
149 {
150 *left = fq.fma_limb(*left, *right, c.clone());
151 }
152 for limb in self.limbs_mut()[min_limb..].iter_mut() {
153 *limb = fq.reduce(*limb);
154 }
155 }
156 }
157
158 pub fn add(&mut self, other: &Self, c: FieldElement<F>) {
159 self.add_offset(other, c, 0);
160 }
161
162 pub fn assign(&mut self, other: &Self) {
163 assert_eq!(self.fq(), other.fq());
164 assert_eq!(self.len(), other.len());
165 self.limbs_mut().copy_from_slice(other.limbs())
166 }
167
168 pub fn assign_partial(&mut self, other: &Self) {
170 assert_eq!(self.fq(), other.fq());
171 assert!(other.len() <= self.len());
172
173 self.limbs_mut()[0..other.limbs().len()].copy_from_slice(other.limbs());
174 for limb in self.limbs_mut()[other.limbs().len()..].iter_mut() {
175 *limb = 0;
176 }
177 }
178
179 pub fn is_zero(&self) -> bool {
180 self.limbs().iter().all(|&x| x == 0)
181 }
182
183 pub fn extend_len(&mut self, len: usize) {
186 if self.len() >= len {
187 return;
188 }
189 *self.len_mut() = len;
190 let new_len = self.fq().number(len);
191 self.vec_mut().resize(new_len, 0);
192 }
193
194 pub fn set_scratch_vector_size(&mut self, len: usize) {
197 self.vec_mut().clear();
198 let new_len = self.fq().number(len);
199 self.vec_mut().resize(new_len, 0);
200 *self.len_mut() = len;
201 }
202
203 pub fn copy_from_slice(&mut self, slice: &[FieldElement<F>]) {
206 assert!(slice.iter().all(|x| x.field() == self.fq()));
207 assert_eq!(self.len(), slice.len());
208
209 let fq = self.fq();
210 self.vec_mut().clear();
211 self.vec_mut().extend(
212 slice
213 .chunks(fq.entries_per_limb())
214 .map(|x| fq.pack(x.iter().cloned())),
215 );
216 }
217
218 pub fn sign_rule(&self, other: &Self) -> bool {
219 assert_eq!(self.fq(), other.fq());
220 assert_eq!(self.fq().q(), 2);
221
222 let mut result = 0;
223 for target_limb_idx in 0..self.limbs().len() {
224 let target_limb = other.limbs()[target_limb_idx];
225 let source_limb = self.limbs()[target_limb_idx];
226 result ^= crate::limb::sign_rule(target_limb, source_limb);
227 if target_limb.count_ones().is_multiple_of(2) {
228 continue;
229 }
230 for _ in 0..target_limb_idx {
231 result ^= source_limb.count_ones() % 2;
232 }
233 }
234 result == 1
235 }
236
237 pub fn add_truncate(&mut self, other: &Self, c: FieldElement<F>) -> Option<()> {
238 assert_eq!(self.fq(), other.fq());
239 let fq = self.fq();
240 for (left, right) in self.limbs_mut().iter_mut().zip_eq(other.limbs()) {
241 *left = fq.fma_limb(*left, *right, c.clone());
242 *left = fq.truncate(*left)?;
243 }
244 Some(())
245 }
246
247 fn add_carry_limb<T>(
248 &mut self,
249 idx: usize,
250 source: Limb,
251 c: FieldElement<F>,
252 rest: &mut [T],
253 ) -> bool
254 where
255 for<'a> &'a mut T: TryInto<&'a mut Self>,
256 {
257 assert_eq!(self.fq(), c.field());
258 if self.fq().q() == 2 {
259 let c = self.fq().encode(c);
260 if c == 0 {
261 return false;
262 }
263 let mut cur_vec = self;
264 let mut carry = source;
265 for carry_vec in rest.iter_mut() {
266 let carry_vec = carry_vec
267 .try_into()
268 .ok()
269 .expect("rest vectors in add_carry must be of the same prime");
270 let rem = cur_vec.limbs()[idx] ^ carry;
271 let quot = cur_vec.limbs()[idx] & carry;
272 cur_vec.limbs_mut()[idx] = rem;
273 carry = quot;
274 cur_vec = carry_vec;
275 if quot == 0 {
276 return false;
277 }
278 }
279 cur_vec.limbs_mut()[idx] ^= carry;
280 true
281 } else {
282 unimplemented!()
283 }
284 }
285
286 pub fn add_carry<T>(&mut self, other: &Self, c: FieldElement<F>, rest: &mut [T]) -> bool
287 where
288 for<'a> &'a mut T: TryInto<&'a mut Self>,
289 {
290 assert_eq!(self.fq(), other.fq());
291 let mut result = false;
292 for i in 0..self.limbs().len() {
293 result |= self.add_carry_limb(i, other.limbs()[i], c.clone(), rest);
294 }
295 result
296 }
297
298 pub fn first_nonzero(&self) -> Option<(usize, FieldElement<F>)> {
300 let entries_per_limb = self.fq().entries_per_limb();
301 let bit_length = self.fq().bit_length();
302 let bitmask = self.fq().bitmask();
303 for (i, &limb) in self.limbs().iter().enumerate() {
304 if limb == 0 {
305 continue;
306 }
307 let index = limb.trailing_zeros() as usize / bit_length;
308 return Some((
309 i * entries_per_limb + index,
310 self.fq().decode((limb >> (index * bit_length)) & bitmask),
311 ));
312 }
313 None
314 }
315
316 pub fn density(&self) -> f32 {
317 let num_nonzero = if self.fq().q() == 2 {
318 self.limbs()
319 .iter()
320 .copied()
321 .map(Limb::count_ones)
322 .sum::<u32>() as usize
323 } else {
324 self.iter_nonzero().count()
325 };
326 num_nonzero as f32 / self.len() as f32
327 }
328}
329
330impl<T: AsRef<[FieldElement<F>]>, F: Field> From<(F, T)> for FqVector<F> {
331 fn from(data: (F, T)) -> Self {
332 let (fq, slice) = data;
333 assert!(slice.as_ref().iter().all(|x| x.field() == fq));
334 let mut v = Self::new(fq, slice.as_ref().len());
335 v.copy_from_slice(slice.as_ref());
336 v
337 }
338}
339
340impl<F: Field> From<&FqVector<F>> for Vec<FieldElement<F>> {
341 fn from(vec: &FqVector<F>) -> Self {
342 vec.iter().collect()
343 }
344}
345
346impl<F: Field> std::fmt::Display for FqVector<F> {
347 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
348 self.as_slice().fmt(f)
349 }
350}
351
352#[cfg(feature = "proptest")]
353pub mod arbitrary {
354 use proptest::prelude::*;
355
356 use super::*;
357
358 pub const MAX_LEN: usize = 10_000;
359
360 #[derive(Debug, Clone)]
361 pub struct FqVectorArbParams<F> {
362 pub fq: Option<F>,
363 pub len: BoxedStrategy<usize>,
364 }
365
366 impl<F> Default for FqVectorArbParams<F> {
367 fn default() -> Self {
368 Self {
369 fq: None,
370 len: (0..=MAX_LEN).boxed(),
371 }
372 }
373 }
374
375 impl<F: Field> Arbitrary for FqVector<F> {
376 type Parameters = FqVectorArbParams<F>;
377 type Strategy = BoxedStrategy<Self>;
378
379 fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
380 let fq = match args.fq {
381 Some(fq) => Just(fq).boxed(),
382 None => any::<F>().boxed(),
383 };
384 (fq, args.len)
385 .prop_flat_map(|(fq, len)| {
386 (Just(fq), proptest::collection::vec(fq.arb_element(), len))
387 })
388 .prop_map(|(fq, v)| Self::from_slice(fq, &v))
389 .boxed()
390 }
391 }
392}