1use super::inner::FqSlice;
2use crate::{
3 field::{Field, element::FieldElement},
4 limb::Limb,
5};
6
7pub struct FqVectorIterator<'a, F> {
8 fq: F,
9 limbs: &'a [Limb],
10 bit_length: usize,
11 bit_mask: Limb,
12 entries_per_limb_m_1: usize,
13 limb_index: usize,
14 entries_left: usize,
15 cur_limb: Limb,
16 counter: usize,
17}
18
19impl<'a, F: Field> FqVectorIterator<'a, F> {
20 pub(super) fn new(vec: FqSlice<'a, F>) -> Self {
21 let counter = vec.len();
22 let limbs = vec.into_limbs();
23
24 if counter == 0 {
25 return Self {
26 fq: vec.fq(),
27 limbs,
28 bit_length: 0,
29 entries_per_limb_m_1: 0,
30 bit_mask: 0,
31 limb_index: 0,
32 entries_left: 0,
33 cur_limb: 0,
34 counter,
35 };
36 }
37 let pair = vec.fq().limb_bit_index_pair(vec.start());
38
39 let bit_length = vec.fq().bit_length();
40 let cur_limb = limbs[pair.limb] >> pair.bit_index;
41
42 let entries_per_limb = vec.fq().entries_per_limb();
43 Self {
44 fq: vec.fq(),
45 limbs,
46 bit_length,
47 entries_per_limb_m_1: entries_per_limb - 1,
48 bit_mask: vec.fq().bitmask(),
49 limb_index: pair.limb,
50 entries_left: entries_per_limb - (vec.start() % entries_per_limb),
51 cur_limb,
52 counter,
53 }
54 }
55
56 pub fn skip_n(&mut self, mut n: usize) {
57 if n >= self.counter {
58 self.counter = 0;
59 return;
60 }
61 let entries_per_limb = self.entries_per_limb_m_1 + 1;
62 if n < self.entries_left {
63 self.entries_left -= n;
64 self.counter -= n;
65 self.cur_limb >>= self.bit_length * n;
66 return;
67 }
68
69 n -= self.entries_left;
70 self.counter -= self.entries_left;
71 self.entries_left = 0;
72
73 let skip_limbs = n / entries_per_limb;
74 self.limb_index += skip_limbs;
75 self.counter -= skip_limbs * entries_per_limb;
76 n -= skip_limbs * entries_per_limb;
77
78 if n > 0 {
79 self.entries_left = entries_per_limb - n;
80 self.limb_index += 1;
81 self.cur_limb = self.limbs[self.limb_index] >> (n * self.bit_length);
82 self.counter -= n;
83 }
84 }
85}
86
87impl<F: Field> Iterator for FqVectorIterator<'_, F> {
88 type Item = FieldElement<F>;
89
90 fn next(&mut self) -> Option<Self::Item> {
91 if self.counter == 0 {
92 return None;
93 } else if self.entries_left == 0 {
94 self.limb_index += 1;
95 self.cur_limb = self.limbs[self.limb_index];
96 self.entries_left = self.entries_per_limb_m_1;
97 } else {
98 self.entries_left -= 1;
99 }
100
101 let result = self.cur_limb & self.bit_mask;
102 self.counter -= 1;
103 self.cur_limb >>= self.bit_length;
104
105 Some(self.fq.decode(result))
106 }
107}
108
109impl<F: Field> ExactSizeIterator for FqVectorIterator<'_, F> {
110 fn len(&self) -> usize {
111 self.counter
112 }
113}
114
115pub struct FqVectorNonZeroIterator<'a, F> {
118 fq: F,
119 limbs: &'a [Limb],
120 limb_index: usize,
121 cur_limb_entries_left: usize,
122 cur_limb: Limb,
123 idx: usize,
124 dim: usize,
125}
126
127impl<'a, F: Field> FqVectorNonZeroIterator<'a, F> {
128 pub(super) fn new(vec: FqSlice<'a, F>) -> Self {
129 let entries_per_limb = vec.fq().entries_per_limb();
130
131 let dim = vec.len();
132 let limbs = vec.into_limbs();
133
134 if dim == 0 {
135 return Self {
136 fq: vec.fq(),
137 limbs,
138 limb_index: 0,
139 cur_limb_entries_left: 0,
140 cur_limb: 0,
141 idx: 0,
142 dim: 0,
143 };
144 }
145 let min_index = vec.start();
146 let pair = vec.fq().limb_bit_index_pair(min_index);
147 let cur_limb = limbs[pair.limb] >> pair.bit_index;
148 let cur_limb_entries_left = entries_per_limb - (min_index % entries_per_limb);
149 Self {
150 fq: vec.fq(),
151 limbs,
152 limb_index: pair.limb,
153 cur_limb_entries_left,
154 cur_limb,
155 idx: 0,
156 dim,
157 }
158 }
159}
160
161impl<F: Field> Iterator for FqVectorNonZeroIterator<'_, F> {
162 type Item = (usize, FieldElement<F>);
163
164 fn next(&mut self) -> Option<Self::Item> {
165 let bit_length: usize = self.fq.bit_length();
166 let bitmask: Limb = self.fq.bitmask();
167 let entries_per_limb: usize = self.fq.entries_per_limb();
168 loop {
169 let bits_left = (self.cur_limb_entries_left * bit_length) as u32;
170 #[allow(clippy::unnecessary_cast)]
171 let tz_real = (self.cur_limb | (1 as Limb).checked_shl(bits_left as u32).unwrap_or(0))
172 .trailing_zeros();
173 let tz_rem = ((tz_real as u8) % (bit_length as u8)) as u32;
174 let tz_div = ((tz_real as u8) / (bit_length as u8)) as u32;
175 let tz = tz_real - tz_rem;
176 self.idx += tz_div as usize;
177 if self.idx >= self.dim {
178 return None;
179 }
180 self.cur_limb_entries_left -= tz_div as usize;
181 if self.cur_limb_entries_left == 0 {
182 self.limb_index += 1;
183 self.cur_limb_entries_left = entries_per_limb;
184 self.cur_limb = self.limbs[self.limb_index];
185 continue;
186 }
187 self.cur_limb >>= tz;
188 if tz == 0 {
189 break;
190 }
191 }
192 let result = (self.idx, self.fq.decode(self.cur_limb & bitmask));
193 self.idx += 1;
194 self.cur_limb_entries_left -= 1;
195 self.cur_limb >>= bit_length;
196 Some(result)
197 }
198}