Skip to main content

hekate_core/poly/
variant.rs

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