fp/vector/
impl_fqvector.rs

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        // self.limbs is allowed to have more limbs than necessary, but we only save the
49        // necessary ones.
50        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        // This is sound because `fq.encode(fq.zero())` is always zero.
110        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    /// Add `other` to `self` on the assumption that the first `offset` entries of `other` are
130    /// empty.
131    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    /// A version of [`FqVector::assign`] that allows `other` to be shorter than `self`.
169    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    /// This function ensures the length of the vector is at least `len`. See also
184    /// `set_scratch_vector_size`.
185    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    /// This clears the vector and sets the length to `len`. This is useful for reusing
195    /// allocations of temporary vectors.
196    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    /// This replaces the contents of the vector with the contents of the slice. The two must have
204    /// the same length.
205    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    /// Find the index and value of the first non-zero entry of the vector. `None` if the vector is zero.
299    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}