Skip to main content

hekate_core/poly/
variant.rs

1// SPDX-FileCopyrightText: 2026 Andrei Kochergin <andrei@oumuamua.dev>
2// SPDX-FileCopyrightText: 2026 Oumuamua Labs <info@oumuamua.dev>
3// SPDX-License-Identifier: AGPL-3.0-only
4
5use crate::errors;
6use crate::tensor::TensorProduct;
7use crate::trace::TraceCompatibleField;
8use alloc::vec;
9use alloc::vec::Vec;
10use core::fmt;
11use hekate_math::{Bit, Block8, Block16, Block32, Block64, Block128, Flat};
12use zeroize::Zeroize;
13
14/// Failures raised by `PolyVariant` operations.
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
16pub enum Error {
17    /// `1 << num_vars` overflowed `usize`.
18    DomainTooLarge { num_vars: usize },
19
20    /// Polynomial length disagrees with `2^num_vars`.
21    DomainSizeMismatch { expected_len: usize, got_len: usize },
22
23    /// Evaluation point has the
24    /// wrong number of coordinates.
25    PointDimensionMismatch { expected_len: usize, got_len: usize },
26
27    /// Variant is read-only at fold time.
28    UnsupportedFold { kind: &'static str },
29}
30
31impl fmt::Display for Error {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        match self {
34            Self::DomainTooLarge { num_vars } => {
35                write!(
36                    f,
37                    "Virtual polynomial domain too large: num_vars={num_vars}"
38                )
39            }
40            Self::DomainSizeMismatch {
41                expected_len,
42                got_len,
43            } => write!(
44                f,
45                "Virtual polynomial domain size mismatch: expected {expected_len}, got {got_len}",
46            ),
47            Self::PointDimensionMismatch {
48                expected_len,
49                got_len,
50            } => write!(
51                f,
52                "Virtual polynomial point dimension mismatch: expected {expected_len}, got {got_len}",
53            ),
54            Self::UnsupportedFold { kind } => {
55                write!(
56                    f,
57                    "Virtual polynomial cannot be folded lazily for kind: {kind}"
58                )
59            }
60        }
61    }
62}
63
64/// Zero-copy MLE view over a physical trace column.
65/// Every variant must deliver `get_at(i)` without
66/// heap allocation, the hot path inside Sumcheck.
67#[derive(Clone, Debug, Zeroize)]
68pub enum PolyVariant<'a, F>
69where
70    F: TraceCompatibleField,
71{
72    /// Fully materialized hypercube.
73    #[zeroize(skip)]
74    Dense(&'a [Flat<F>]),
75    #[zeroize(skip)]
76    Shifted(&'a [Flat<F>]),
77
78    /// `Eq(x, r)` held lazily as a `TensorProduct`:
79    /// `O(num_vars)` memory, `O(1)` fold.
80    Eq(TensorProduct<F>),
81
82    /// `(data[i] >> bit_idx) & 1` on a `B8` column.
83    #[zeroize(skip)]
84    PackedBitB8 {
85        data: &'a [Flat<Block8>],
86        bit_idx: usize,
87    },
88
89    /// `(data[i] >> bit_idx) & 1` on a `B16` column.
90    #[zeroize(skip)]
91    PackedBitB16 {
92        data: &'a [Flat<Block16>],
93        bit_idx: usize,
94    },
95
96    /// `(data[i] >> bit_idx) & 1` on a `B32` column.
97    #[zeroize(skip)]
98    PackedBitB32 {
99        data: &'a [Flat<Block32>],
100        bit_idx: usize,
101    },
102
103    /// `(data[i] >> bit_idx) & 1` on a `B64` column.
104    #[zeroize(skip)]
105    PackedBitB64 {
106        data: &'a [Flat<Block64>],
107        bit_idx: usize,
108    },
109
110    /// `S(i) = Σ_k cols[k][i]` over boolean
111    /// columns (one-hot selector groups).
112    #[zeroize(skip)]
113    CompositeSelector(Vec<&'a [Bit]>),
114
115    /// Mask that is `1` everywhere and
116    /// `1 - product_of_challenges` at
117    /// index `2^N - 1`. Used to kill
118    /// cross-row wrap in Sumcheck.
119    #[zeroize(skip)]
120    TransitionMask {
121        num_vars: usize,
122        product_of_challenges: F,
123    },
124
125    // ==============================================================
126    // Indirect:
127    // P(x) = data[indices[x]].
128    // Zero-copy permutations and arbitrary wiring.
129    // ==============================================================
130    #[zeroize(skip)]
131    IndirectBit {
132        data: &'a [Bit],
133        indices: &'a [usize],
134    },
135    #[zeroize(skip)]
136    IndirectB8 {
137        data: &'a [Flat<Block8>],
138        indices: &'a [usize],
139    },
140    #[zeroize(skip)]
141    IndirectB16 {
142        data: &'a [Flat<Block16>],
143        indices: &'a [usize],
144    },
145    #[zeroize(skip)]
146    IndirectB32 {
147        data: &'a [Flat<Block32>],
148        indices: &'a [usize],
149    },
150    #[zeroize(skip)]
151    IndirectB64 {
152        data: &'a [Flat<Block64>],
153        indices: &'a [usize],
154    },
155    #[zeroize(skip)]
156    IndirectB128 {
157        data: &'a [Flat<Block128>],
158        indices: &'a [usize],
159    },
160
161    // ====================================================
162    // Stride Access:
163    // P(i) = data[start + i * step].
164    // ====================================================
165    #[zeroize(skip)]
166    StrideBit {
167        data: &'a [Bit],
168        start: usize,
169        step: usize,
170        len: usize,
171    },
172    #[zeroize(skip)]
173    StrideB8 {
174        data: &'a [Flat<Block8>],
175        start: usize,
176        step: usize,
177        len: usize,
178    },
179    #[zeroize(skip)]
180    StrideB16 {
181        data: &'a [Flat<Block16>],
182        start: usize,
183        step: usize,
184        len: usize,
185    },
186    #[zeroize(skip)]
187    StrideB32 {
188        data: &'a [Flat<Block32>],
189        start: usize,
190        step: usize,
191        len: usize,
192    },
193    #[zeroize(skip)]
194    StrideB64 {
195        data: &'a [Flat<Block64>],
196        start: usize,
197        step: usize,
198        len: usize,
199    },
200    #[zeroize(skip)]
201    StrideB128 {
202        data: &'a [Flat<Block128>],
203        start: usize,
204        step: usize,
205        len: usize,
206    },
207
208    // ====================================================
209    // Cyclic Rotation:
210    // P(i) = data[(i + rotation) % len].
211    // Uses bitwise masking for modulo (len must be power of 2).
212    // ====================================================
213    #[zeroize(skip)]
214    RotationBit { data: &'a [Bit], rotation: usize },
215    #[zeroize(skip)]
216    RotationB8 {
217        data: &'a [Flat<Block8>],
218        rotation: usize,
219    },
220    #[zeroize(skip)]
221    RotationB16 {
222        data: &'a [Flat<Block16>],
223        rotation: usize,
224    },
225    #[zeroize(skip)]
226    RotationB32 {
227        data: &'a [Flat<Block32>],
228        rotation: usize,
229    },
230    #[zeroize(skip)]
231    RotationB64 {
232        data: &'a [Flat<Block64>],
233        rotation: usize,
234    },
235    #[zeroize(skip)]
236    RotationB128 {
237        data: &'a [Flat<Block128>],
238        rotation: usize,
239    },
240
241    // ====================================================
242    // Compressed slice views (JIT-promoted to F)
243    // ====================================================
244    #[zeroize(skip)]
245    BitSlice(&'a [Bit]),
246    #[zeroize(skip)]
247    B8Slice(&'a [Flat<Block8>]),
248    #[zeroize(skip)]
249    B16Slice(&'a [Flat<Block16>]),
250    #[zeroize(skip)]
251    B32Slice(&'a [Flat<Block32>]),
252    #[zeroize(skip)]
253    B64Slice(&'a [Flat<Block64>]),
254    #[zeroize(skip)]
255    B128Slice(&'a [Flat<Block128>]),
256
257    // ====================================================
258    // Same-width views shifted by one row (cyclic)
259    // ====================================================
260    #[zeroize(skip)]
261    ShiftedBitSlice(&'a [Bit]),
262    #[zeroize(skip)]
263    ShiftedB8Slice(&'a [Flat<Block8>]),
264    #[zeroize(skip)]
265    ShiftedB16Slice(&'a [Flat<Block16>]),
266    #[zeroize(skip)]
267    ShiftedB32Slice(&'a [Flat<Block32>]),
268    #[zeroize(skip)]
269    ShiftedB64Slice(&'a [Flat<Block64>]),
270    #[zeroize(skip)]
271    ShiftedB128Slice(&'a [Flat<Block128>]),
272
273    #[zeroize(skip)]
274    ShiftedPackedBitB8 {
275        data: &'a [Flat<Block8>],
276        bit_idx: usize,
277    },
278    #[zeroize(skip)]
279    ShiftedPackedBitB16 {
280        data: &'a [Flat<Block16>],
281        bit_idx: usize,
282    },
283    #[zeroize(skip)]
284    ShiftedPackedBitB32 {
285        data: &'a [Flat<Block32>],
286        bit_idx: usize,
287    },
288    #[zeroize(skip)]
289    ShiftedPackedBitB64 {
290        data: &'a [Flat<Block64>],
291        bit_idx: usize,
292    },
293}
294
295impl<'a, F> PolyVariant<'a, F>
296where
297    F: TraceCompatibleField,
298{
299    /// Number of hypercube points (polynomial length).
300    pub fn len(&self) -> usize {
301        match self {
302            Self::Dense(h) => h.len(),
303            Self::Shifted(h) => h.len(),
304            Self::Eq(t) => 1 << t.num_vars(),
305            Self::PackedBitB8 { data, .. } => data.len(),
306            Self::PackedBitB16 { data, .. } => data.len(),
307            Self::PackedBitB32 { data, .. } => data.len(),
308            Self::PackedBitB64 { data, .. } => data.len(),
309            Self::TransitionMask { num_vars, .. } => 1 << num_vars,
310            Self::CompositeSelector(cols) => {
311                if cols.is_empty() {
312                    0
313                } else {
314                    cols[0].len()
315                }
316            }
317            Self::IndirectBit { indices, .. } => indices.len(),
318            Self::IndirectB8 { indices, .. } => indices.len(),
319            Self::IndirectB16 { indices, .. } => indices.len(),
320            Self::IndirectB32 { indices, .. } => indices.len(),
321            Self::IndirectB64 { indices, .. } => indices.len(),
322            Self::IndirectB128 { indices, .. } => indices.len(),
323            Self::StrideBit { len, .. } => *len,
324            Self::StrideB8 { len, .. } => *len,
325            Self::StrideB16 { len, .. } => *len,
326            Self::StrideB32 { len, .. } => *len,
327            Self::StrideB64 { len, .. } => *len,
328            Self::StrideB128 { len, .. } => *len,
329            Self::RotationBit { data, .. } => data.len(),
330            Self::RotationB8 { data, .. } => data.len(),
331            Self::RotationB16 { data, .. } => data.len(),
332            Self::RotationB32 { data, .. } => data.len(),
333            Self::RotationB64 { data, .. } => data.len(),
334            Self::RotationB128 { data, .. } => data.len(),
335            Self::BitSlice(h) => h.len(),
336            Self::B8Slice(h) => h.len(),
337            Self::B16Slice(h) => h.len(),
338            Self::B32Slice(h) => h.len(),
339            Self::B64Slice(h) => h.len(),
340            Self::B128Slice(h) => h.len(),
341            Self::ShiftedBitSlice(h) => h.len(),
342            Self::ShiftedB8Slice(h) => h.len(),
343            Self::ShiftedB16Slice(h) => h.len(),
344            Self::ShiftedB32Slice(h) => h.len(),
345            Self::ShiftedB64Slice(h) => h.len(),
346            Self::ShiftedB128Slice(h) => h.len(),
347            Self::ShiftedPackedBitB8 { data, .. } => data.len(),
348            Self::ShiftedPackedBitB16 { data, .. } => data.len(),
349            Self::ShiftedPackedBitB32 { data, .. } => data.len(),
350            Self::ShiftedPackedBitB64 { data, .. } => data.len(),
351        }
352    }
353
354    pub fn is_empty(&self) -> bool {
355        self.len() == 0
356    }
357
358    /// Read the hypercube value at `index`,
359    /// lifted into `Flat<F>`. `O(1)` for
360    /// slice variants, `O(num_vars)` for `Eq`.
361    #[inline(always)]
362    pub fn get_at(&self, index: usize) -> Flat<F> {
363        match self {
364            Self::Dense(h) => h[index],
365            Self::Shifted(h) => {
366                let len = h.len();
367                let next_idx = if index + 1 == len { 0 } else { index + 1 };
368
369                h[next_idx]
370            }
371            Self::Eq(t) => t.evaluate_at_index(index),
372            Self::PackedBitB8 { data, bit_idx } => {
373                let bit = data[index].tower_bit(*bit_idx);
374                if bit == 1 {
375                    Flat::from_raw(F::ONE)
376                } else {
377                    Flat::from_raw(F::ZERO)
378                }
379            }
380            Self::PackedBitB16 { data, bit_idx } => {
381                let bit = data[index].tower_bit(*bit_idx);
382                if bit == 1 {
383                    Flat::from_raw(F::ONE)
384                } else {
385                    Flat::from_raw(F::ZERO)
386                }
387            }
388            Self::PackedBitB32 { data, bit_idx } => {
389                let bit = data[index].tower_bit(*bit_idx);
390                if bit == 1 {
391                    Flat::from_raw(F::ONE)
392                } else {
393                    Flat::from_raw(F::ZERO)
394                }
395            }
396            Self::PackedBitB64 { data, bit_idx } => {
397                let bit = data[index].tower_bit(*bit_idx);
398                if bit == 1 {
399                    Flat::from_raw(F::ONE)
400                } else {
401                    Flat::from_raw(F::ZERO)
402                }
403            }
404            Self::TransitionMask {
405                num_vars,
406                product_of_challenges,
407            } => {
408                let last_idx: usize = (1 << num_vars) - 1;
409                if index == last_idx {
410                    Flat::from_raw(F::ONE - *product_of_challenges)
411                } else {
412                    Flat::from_raw(F::ONE)
413                }
414            }
415            Self::CompositeSelector(cols) => {
416                let mut sum = Flat::from_raw(F::default());
417                for col in cols {
418                    if col[index].get() == 1 {
419                        sum += Flat::from_raw(F::ONE);
420                    }
421                }
422
423                sum
424            }
425
426            Self::IndirectBit { data, indices } => Flat::from_raw(F::from(data[indices[index]])),
427            Self::IndirectB8 { data, indices } => F::promote_flat(data[indices[index]]),
428            Self::IndirectB16 { data, indices } => F::promote_flat(data[indices[index]]),
429            Self::IndirectB32 { data, indices } => F::promote_flat(data[indices[index]]),
430            Self::IndirectB64 { data, indices } => F::promote_flat(data[indices[index]]),
431            Self::IndirectB128 { data, indices } => F::promote_flat(data[indices[index]]),
432
433            Self::StrideBit {
434                data, start, step, ..
435            } => Flat::from_raw(F::from(data[start + index * step])),
436            Self::StrideB8 {
437                data, start, step, ..
438            } => F::promote_flat(data[start + index * step]),
439            Self::StrideB16 {
440                data, start, step, ..
441            } => F::promote_flat(data[start + index * step]),
442            Self::StrideB32 {
443                data, start, step, ..
444            } => F::promote_flat(data[start + index * step]),
445            Self::StrideB64 {
446                data, start, step, ..
447            } => F::promote_flat(data[start + index * step]),
448            Self::StrideB128 {
449                data, start, step, ..
450            } => F::promote_flat(data[start + index * step]),
451
452            Self::RotationBit { data, rotation } => {
453                Flat::from_raw(F::from(data[(index + rotation) & (data.len() - 1)]))
454            }
455            Self::RotationB8 { data, rotation } => {
456                F::promote_flat(data[(index + rotation) & (data.len() - 1)])
457            }
458            Self::RotationB16 { data, rotation } => {
459                F::promote_flat(data[(index + rotation) & (data.len() - 1)])
460            }
461            Self::RotationB32 { data, rotation } => {
462                F::promote_flat(data[(index + rotation) & (data.len() - 1)])
463            }
464            Self::RotationB64 { data, rotation } => {
465                F::promote_flat(data[(index + rotation) & (data.len() - 1)])
466            }
467            Self::RotationB128 { data, rotation } => {
468                F::promote_flat(data[(index + rotation) & (data.len() - 1)])
469            }
470
471            Self::BitSlice(s) => Flat::from_raw(F::from(s[index])),
472            Self::B8Slice(s) => F::promote_flat(s[index]),
473            Self::B16Slice(s) => F::promote_flat(s[index]),
474            Self::B32Slice(s) => F::promote_flat(s[index]),
475            Self::B64Slice(s) => F::promote_flat(s[index]),
476            Self::B128Slice(s) => F::promote_flat(s[index]),
477
478            Self::ShiftedBitSlice(s) => {
479                let len = s.len();
480                let next_idx = if index + 1 == len { 0 } else { index + 1 };
481
482                Flat::from_raw(F::from(s[next_idx]))
483            }
484            Self::ShiftedB8Slice(s) => {
485                let len = s.len();
486                let next_idx = if index + 1 == len { 0 } else { index + 1 };
487
488                F::promote_flat(s[next_idx])
489            }
490            Self::ShiftedB16Slice(s) => {
491                let len = s.len();
492                let next_idx = if index + 1 == len { 0 } else { index + 1 };
493
494                F::promote_flat(s[next_idx])
495            }
496            Self::ShiftedB32Slice(s) => {
497                let len = s.len();
498                let next_idx = if index + 1 == len { 0 } else { index + 1 };
499
500                F::promote_flat(s[next_idx])
501            }
502            Self::ShiftedB64Slice(s) => {
503                let len = s.len();
504                let next_idx = if index + 1 == len { 0 } else { index + 1 };
505
506                F::promote_flat(s[next_idx])
507            }
508            Self::ShiftedB128Slice(s) => {
509                let len = s.len();
510                let next_idx = if index + 1 == len { 0 } else { index + 1 };
511
512                F::promote_flat(s[next_idx])
513            }
514            Self::ShiftedPackedBitB8 { data, bit_idx } => {
515                let len = data.len();
516                let next_idx = if index + 1 == len { 0 } else { index + 1 };
517                let bit = data[next_idx].tower_bit(*bit_idx);
518
519                if bit == 1 {
520                    Flat::from_raw(F::ONE)
521                } else {
522                    Flat::from_raw(F::ZERO)
523                }
524            }
525            Self::ShiftedPackedBitB16 { data, bit_idx } => {
526                let len = data.len();
527                let next_idx = if index + 1 == len { 0 } else { index + 1 };
528                let bit = data[next_idx].tower_bit(*bit_idx);
529
530                if bit == 1 {
531                    Flat::from_raw(F::ONE)
532                } else {
533                    Flat::from_raw(F::ZERO)
534                }
535            }
536            Self::ShiftedPackedBitB32 { data, bit_idx } => {
537                let len = data.len();
538                let next_idx = if index + 1 == len { 0 } else { index + 1 };
539                let bit = data[next_idx].tower_bit(*bit_idx);
540
541                if bit == 1 {
542                    Flat::from_raw(F::ONE)
543                } else {
544                    Flat::from_raw(F::ZERO)
545                }
546            }
547            Self::ShiftedPackedBitB64 { data, bit_idx } => {
548                let len = data.len();
549                let next_idx = if index + 1 == len { 0 } else { index + 1 };
550                let bit = data[next_idx].tower_bit(*bit_idx);
551
552                if bit == 1 {
553                    Flat::from_raw(F::ONE)
554                } else {
555                    Flat::from_raw(F::ZERO)
556                }
557            }
558        }
559    }
560
561    /// Evaluate the MLE at an arbitrary point.
562    /// `Eq` delegates to `TensorProduct`;
563    /// every other variant expands MLE weights
564    /// once and does a single `get_at` sweep.
565    #[inline(always)]
566    pub fn evaluate(&self, point: &[Flat<F>]) -> errors::Result<Flat<F>> {
567        match self {
568            Self::Eq(t) => Ok(t.evaluate_extension(point)?),
569            _ => {
570                let num_vars = point.len();
571                let got_len = self.len();
572
573                let Some(expected_len) = 1usize.checked_shl(num_vars as u32) else {
574                    return Err(Error::DomainTooLarge { num_vars }.into());
575                };
576
577                if got_len != expected_len {
578                    return Err(Error::DomainSizeMismatch {
579                        expected_len,
580                        got_len,
581                    }
582                    .into());
583                }
584
585                if num_vars == 0 {
586                    return Ok(self.get_at(0));
587                }
588
589                let weights = Self::expand_mle_weights(point);
590                let mut total = Flat::from_raw(F::ZERO);
591
592                for (i, w) in weights.iter().enumerate() {
593                    total += self.get_at(i) * *w;
594                }
595
596                Ok(total)
597            }
598        }
599    }
600
601    pub fn expand_mle_weights(r: &[Flat<F>]) -> Vec<Flat<F>> {
602        let num_vars = r.len();
603        let size = 1 << num_vars;
604
605        let mut weights = vec![Flat::from_raw(F::ZERO); size];
606        weights[0] = Flat::from_raw(F::ONE);
607
608        for (i, &rk) in r.iter().enumerate() {
609            let one_minus_rk = Flat::from_raw(F::ONE) - rk;
610            let current_len = 1 << i;
611
612            for i in 0..current_len {
613                let w = weights[i];
614                weights[i] = w * one_minus_rk;
615                weights[current_len + i] = w * rk;
616            }
617        }
618
619        weights
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626    use hekate_math::{Bit, Block8, Block128, FlatPromote, HardwareField, TowerField};
627
628    type F = Block128;
629
630    #[test]
631    fn get_at_dense_returns_correct_values() {
632        let data: Vec<Flat<F>> = (0..4u128).map(|i| F::from(i * 10).to_hardware()).collect();
633        let v = PolyVariant::<F>::Dense(&data);
634
635        assert_eq!(v.get_at(0), F::from(0u128).to_hardware());
636        assert_eq!(v.get_at(1), F::from(10u128).to_hardware());
637        assert_eq!(v.get_at(2), F::from(20u128).to_hardware());
638        assert_eq!(v.get_at(3), F::from(30u128).to_hardware());
639    }
640
641    #[test]
642    fn get_at_bit_slice_promotes_to_field() {
643        let bits = vec![
644            Bit::from(0u32),
645            Bit::from(1u32),
646            Bit::from(1u32),
647            Bit::from(0u32),
648        ];
649        let v = PolyVariant::<F>::BitSlice(&bits);
650
651        assert_eq!(v.get_at(0), Flat::from_raw(F::ZERO));
652        assert_eq!(v.get_at(1), Flat::from_raw(F::ONE));
653        assert_eq!(v.get_at(2), Flat::from_raw(F::ONE));
654        assert_eq!(v.get_at(3), Flat::from_raw(F::ZERO));
655    }
656
657    #[test]
658    fn get_at_b8_slice_promotes() {
659        let data = vec![
660            Block8::from(0u8).to_hardware(),
661            Block8::from(0xFFu8).to_hardware(),
662        ];
663        let v = PolyVariant::<F>::B8Slice(&data);
664
665        assert_eq!(
666            v.get_at(0),
667            F::promote_flat(Block8::from(0u8).to_hardware())
668        );
669        assert_eq!(
670            v.get_at(1),
671            F::promote_flat(Block8::from(0xFFu8).to_hardware())
672        );
673    }
674
675    #[test]
676    fn len_matches_data_size() {
677        let data = vec![Flat::from_raw(F::ZERO); 16];
678        assert_eq!(PolyVariant::<F>::Dense(&data).len(), 16);
679
680        let bits = vec![Bit::from(0u32); 8];
681        assert_eq!(PolyVariant::<F>::BitSlice(&bits).len(), 8);
682
683        let eq = TensorProduct::new(vec![Flat::from_raw(F::ONE); 5]);
684        assert_eq!(PolyVariant::<F>::Eq(eq).len(), 32);
685
686        let empty: Vec<Flat<F>> = vec![];
687        assert!(PolyVariant::<F>::Dense(&empty).is_empty());
688    }
689
690    #[test]
691    fn evaluate_constant_polynomial() {
692        let num_vars = 3;
693        let data = vec![F::from(42u128).to_hardware(); 1 << num_vars];
694        let v = PolyVariant::<F>::Dense(&data);
695
696        let point: Vec<Flat<F>> = vec![
697            Flat::from_raw(F::from(1u128).to_hardware().into_raw()),
698            Flat::from_raw(F::from(2u128).to_hardware().into_raw()),
699            Flat::from_raw(F::from(3u128).to_hardware().into_raw()),
700        ];
701
702        let val = v.evaluate(&point).unwrap();
703        assert_eq!(val.into_raw(), F::from(42u128).to_hardware().into_raw());
704    }
705
706    #[test]
707    fn evaluate_linear_polynomial() {
708        let data = vec![F::ZERO.to_hardware(), F::from(10u128).to_hardware()];
709        let v = PolyVariant::<F>::Dense(&data);
710
711        let point = vec![Flat::from_raw(F::from(2u128).to_hardware().into_raw())];
712        let val = v.evaluate(&point).unwrap();
713        assert_eq!(val.into_raw(), F::from(20u128).to_hardware().into_raw());
714    }
715
716    #[test]
717    fn evaluate_single_row() {
718        let data = vec![F::from(99u128).to_hardware()];
719        let v = PolyVariant::<F>::Dense(&data);
720
721        let val = v.evaluate(&[]).unwrap();
722        assert_eq!(val.into_raw(), F::from(99u128).to_hardware().into_raw());
723    }
724
725    #[test]
726    fn evaluate_domain_mismatch_rejected() {
727        let data = vec![F::ZERO.to_hardware(); 4];
728        let v = PolyVariant::<F>::Dense(&data);
729
730        let point = vec![Flat::from_raw(F::ONE); 3];
731        assert!(v.evaluate(&point).is_err());
732    }
733
734    #[test]
735    fn evaluate_eq_polynomial() {
736        let r = vec![Flat::from_raw(F::ONE), Flat::from_raw(F::ZERO)];
737        let eq = PolyVariant::<F>::Eq(TensorProduct::new(r.clone()));
738
739        let val = eq.evaluate(&r).unwrap();
740        assert_eq!(val.into_raw(), F::ONE.to_hardware().into_raw());
741    }
742
743    #[test]
744    fn expand_mle_weights_single_var() {
745        let r = vec![Flat::from_raw(F::from(7u128).to_hardware().into_raw())];
746        let w = PolyVariant::<F>::expand_mle_weights(&r);
747
748        assert_eq!(w.len(), 2);
749        let r0 = r[0];
750        let one = Flat::from_raw(F::ONE);
751        assert_eq!(w[0], one - r0);
752        assert_eq!(w[1], r0);
753    }
754
755    #[test]
756    fn expand_mle_weights_zero_vars() {
757        let w = PolyVariant::<F>::expand_mle_weights(&[]);
758        assert_eq!(w.len(), 1);
759        assert_eq!(w[0], Flat::from_raw(F::ONE));
760    }
761
762    #[test]
763    fn shifted_get_at_wraps_cyclically() {
764        let data: Vec<Flat<F>> = (1..=4u128).map(|i| F::from(i).to_hardware()).collect();
765        let v = PolyVariant::<F>::Shifted(&data);
766
767        assert_eq!(v.get_at(0), F::from(2u128).to_hardware());
768        assert_eq!(v.get_at(1), F::from(3u128).to_hardware());
769        assert_eq!(v.get_at(2), F::from(4u128).to_hardware());
770        assert_eq!(v.get_at(3), F::from(1u128).to_hardware());
771    }
772
773    #[test]
774    fn composite_selector_sums_columns() {
775        let a = vec![
776            Bit::from(1u32),
777            Bit::from(0u32),
778            Bit::from(1u32),
779            Bit::from(0u32),
780        ];
781        let b = vec![
782            Bit::from(0u32),
783            Bit::from(1u32),
784            Bit::from(1u32),
785            Bit::from(0u32),
786        ];
787        let v = PolyVariant::<F>::CompositeSelector(vec![&a, &b]);
788
789        assert_eq!(v.get_at(0), Flat::from_raw(F::ONE));
790        assert_eq!(v.get_at(1), Flat::from_raw(F::ONE));
791
792        let two = F::ONE + F::ONE;
793        assert_eq!(v.get_at(2), Flat::from_raw(two));
794        assert_eq!(v.get_at(3), Flat::from_raw(F::ZERO));
795    }
796}