From 8f21fddd04a179f9f62332466f08ddb8c44843ee Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 6 Feb 2022 23:06:04 -0800
Subject: [PATCH] Add a `PolynomialValues::selector` method for convenience
 (#470)

Also adds a test for `eval_l_1_and_l_last`.
---
 field/src/polynomial/mod.rs | 11 +++++++++++
 starky/src/prover.rs        | 13 +++----------
 starky/src/stark_testing.rs | 12 ++----------
 starky/src/verifier.rs      | 24 ++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/field/src/polynomial/mod.rs b/field/src/polynomial/mod.rs
index 7fccb46e..4d6d55c2 100644
--- a/field/src/polynomial/mod.rs
+++ b/field/src/polynomial/mod.rs
@@ -26,6 +26,17 @@ impl<F: Field> PolynomialValues<F> {
         PolynomialValues { values }
     }
 
+    pub fn zero(len: usize) -> Self {
+        Self::new(vec![F::ZERO; len])
+    }
+
+    /// Returns the polynomial whole value is one at the given index, and zero elsewhere.
+    pub fn selector(len: usize, index: usize) -> Self {
+        let mut result = Self::zero(len);
+        result.values[index] = F::ONE;
+        result
+    }
+
     /// The number of values stored.
     pub fn len(&self) -> usize {
         self.values.len()
diff --git a/starky/src/prover.rs b/starky/src/prover.rs
index d8913b9d..35c30cbf 100644
--- a/starky/src/prover.rs
+++ b/starky/src/prover.rs
@@ -172,17 +172,10 @@ where
     let next_step = 1 << quotient_degree_bits;
 
     // Evaluation of the first Lagrange polynomial on the LDE domain.
-    let lagrange_first = {
-        let mut evals = PolynomialValues::new(vec![F::ZERO; degree]);
-        evals.values[0] = F::ONE;
-        evals.lde_onto_coset(quotient_degree_bits)
-    };
+    let lagrange_first = PolynomialValues::selector(degree, 0).lde_onto_coset(quotient_degree_bits);
     // Evaluation of the last Lagrange polynomial on the LDE domain.
-    let lagrange_last = {
-        let mut evals = PolynomialValues::new(vec![F::ZERO; degree]);
-        evals.values[degree - 1] = F::ONE;
-        evals.lde_onto_coset(quotient_degree_bits)
-    };
+    let lagrange_last =
+        PolynomialValues::selector(degree, degree - 1).lde_onto_coset(quotient_degree_bits);
 
     let z_h_on_coset = ZeroPolyOnCoset::<F>::new(degree_bits, quotient_degree_bits);
 
diff --git a/starky/src/stark_testing.rs b/starky/src/stark_testing.rs
index 8d41e645..222ebf39 100644
--- a/starky/src/stark_testing.rs
+++ b/starky/src/stark_testing.rs
@@ -27,16 +27,8 @@ where
     let size = trace_ldes.len();
     let public_inputs = F::rand_arr::<{ S::PUBLIC_INPUTS }>();
 
-    let lagrange_first = {
-        let mut evals = PolynomialValues::new(vec![F::ZERO; WITNESS_SIZE]);
-        evals.values[0] = F::ONE;
-        evals.lde(rate_bits)
-    };
-    let lagrange_last = {
-        let mut evals = PolynomialValues::new(vec![F::ZERO; WITNESS_SIZE]);
-        evals.values[WITNESS_SIZE - 1] = F::ONE;
-        evals.lde(rate_bits)
-    };
+    let lagrange_first = PolynomialValues::selector(WITNESS_SIZE, 0).lde(rate_bits);
+    let lagrange_last = PolynomialValues::selector(WITNESS_SIZE, WITNESS_SIZE - 1).lde(rate_bits);
 
     let last = F::primitive_root_of_unity(log2_strict(WITNESS_SIZE)).inverse();
     let subgroup =
diff --git a/starky/src/verifier.rs b/starky/src/verifier.rs
index bb0634f5..91a51bed 100644
--- a/starky/src/verifier.rs
+++ b/starky/src/verifier.rs
@@ -148,3 +148,27 @@ fn recover_degree<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, cons
     let lde_bits = config.fri_config.cap_height + initial_merkle_proof.siblings.len();
     1 << (lde_bits - config.fri_config.rate_bits)
 }
+
+#[cfg(test)]
+mod tests {
+    use plonky2::field::field_types::Field;
+    use plonky2::field::goldilocks_field::GoldilocksField;
+    use plonky2::field::polynomial::PolynomialValues;
+
+    use crate::verifier::eval_l_1_and_l_last;
+
+    #[test]
+    fn test_eval_l_1_and_l_last() {
+        type F = GoldilocksField;
+        let log_n = 5;
+        let n = 1 << log_n;
+
+        let x = F::rand(); // challenge point
+        let expected_l_first_x = PolynomialValues::selector(n, 0).ifft().eval(x);
+        let expected_l_last_x = PolynomialValues::selector(n, n - 1).ifft().eval(x);
+
+        let (l_first_x, l_last_x) = eval_l_1_and_l_last(log_n, x);
+        assert_eq!(l_first_x, expected_l_first_x);
+        assert_eq!(l_last_x, expected_l_last_x);
+    }
+}