halo2: Add utilities gadget from orchard

examples/halo2/src/circuit/gadget/utilities.rs

@@ -0,0 +1,299 @@
+use ff::PrimeFieldBits;
+use halo2::{
+    circuit::{Cell, Layouter, Region},
+    plonk::{Advice, Column, Error, Expression},
+};
+use pasta_curves::arithmetic::FieldExt;
+use std::{array, convert::TryInto, ops::Range};
+
+// pub(crate) mod cond_swap;
+// pub(crate) mod decompose_running_sum;
+// pub(crate) mod lookup_range_check;
+
+/// A variable representing a field element.
+#[derive(Copy, Clone, Debug)]
+pub struct CellValue<F: FieldExt> {
+    cell: Cell,
+    value: Option<F>,
+}
+
+pub trait Var<F: FieldExt>: Copy + Clone + std::fmt::Debug {
+    fn new(cell: Cell, value: Option<F>) -> Self;
+    fn cell(&self) -> Cell;
+    fn value(&self) -> Option<F>;
+}
+
+impl<F: FieldExt> Var<F> for CellValue<F> {
+    fn new(cell: Cell, value: Option<F>) -> Self {
+        Self { cell, value }
+    }
+
+    fn cell(&self) -> Cell {
+        self.cell
+    }
+
+    fn value(&self) -> Option<F> {
+        self.value
+    }
+}
+
+pub trait UtilitiesInstructions<F: FieldExt> {
+    type Var: Var<F>;
+
+    fn load_private(
+        &self,
+        mut layouter: impl Layouter<F>,
+        column: Column<Advice>,
+        value: Option<F>,
+    ) -> Result<Self::Var, Error> {
+        layouter.assign_region(
+            || "load private",
+            |mut region| {
+                let cell = region.assign_advice(
+                    || "load private",
+                    column,
+                    0,
+                    || value.ok_or(Error::SynthesisError),
+                )?;
+                Ok(Var::new(cell, value))
+            },
+        )
+    }
+}
+
+/// Assigns a cell at a specific offset within the given region, constraining it
+/// to the same value as another cell (which may be in any region).
+///
+/// Returns an error if either `column` or `copy` is not in a column that was passed to
+/// [`ConstraintSystem::enable_equality`] during circuit configuration.
+///
+/// [`ConstraintSystem::enable_equality`]: halo2::plonk::ConstraintSystem::enable_equality
+pub fn copy<A, AR, F: FieldExt>(
+    region: &mut Region<'_, F>,
+    annotation: A,
+    column: Column<Advice>,
+    offset: usize,
+    copy: &CellValue<F>,
+) -> Result<CellValue<F>, Error>
+where
+    A: Fn() -> AR,
+    AR: Into<String>,
+{
+    let cell = region.assign_advice(annotation, column, offset, || {
+        copy.value.ok_or(Error::SynthesisError)
+    })?;
+
+    region.constrain_equal(cell, copy.cell)?;
+
+    Ok(CellValue::new(cell, copy.value))
+}
+
+pub fn transpose_option_array<T: Copy + std::fmt::Debug, const LEN: usize>(
+    option_array: Option<[T; LEN]>,
+) -> [Option<T>; LEN] {
+    let mut ret = [None; LEN];
+    if let Some(arr) = option_array {
+        for (entry, value) in ret.iter_mut().zip(array::IntoIter::new(arr)) {
+            *entry = Some(value);
+        }
+    }
+    ret
+}
+
+/// Checks that an expresssion is either 1 or 0.
+pub fn bool_check<F: FieldExt>(value: Expression<F>) -> Expression<F> {
+    value.clone() * (Expression::Constant(F::one()) - value)
+}
+
+/// Takes a specified subsequence of the little-endian bit representation of a field element.
+/// The bits are numbered from 0 for the LSB.
+pub fn bitrange_subset<F: FieldExt + PrimeFieldBits>(field_elem: F, bitrange: Range<usize>) -> F {
+    assert!(bitrange.end <= F::NUM_BITS as usize);
+
+    let bits: Vec<bool> = field_elem
+        .to_le_bits()
+        .iter()
+        .by_val()
+        .skip(bitrange.start)
+        .take(bitrange.end - bitrange.start)
+        .chain(std::iter::repeat(false))
+        .take(256)
+        .collect();
+    let bytearray: Vec<u8> = bits
+        .chunks_exact(8)
+        .map(|byte| byte.iter().rev().fold(0u8, |acc, bit| acc * 2 + *bit as u8))
+        .collect();
+
+    F::from_bytes(&bytearray.try_into().unwrap()).unwrap()
+}
+
+/// Check that an expression is in the small range [0..range),
+/// i.e. 0 ≤ word < range.
+pub fn range_check<F: FieldExt>(word: Expression<F>, range: usize) -> Expression<F> {
+    (1..range).fold(word.clone(), |acc, i| {
+        acc * (word.clone() - Expression::Constant(F::from_u64(i as u64)))
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bigint::U256;
+    use ff::PrimeField;
+    use halo2::{
+        circuit::{Layouter, SimpleFloorPlanner},
+        dev::{MockProver, VerifyFailure},
+        plonk::{Circuit, ConstraintSystem, Error, Selector},
+        poly::Rotation,
+    };
+    use pasta_curves::pallas;
+
+    #[test]
+    fn test_range_check() {
+        struct MyCircuit<const RANGE: usize>(u8);
+
+        impl<const RANGE: usize> UtilitiesInstructions<pallas::Base> for MyCircuit<RANGE> {
+            type Var = CellValue<pallas::Base>;
+        }
+
+        #[derive(Clone)]
+        struct Config {
+            selector: Selector,
+            advice: Column<Advice>,
+        }
+
+        impl<const RANGE: usize> Circuit<pallas::Base> for MyCircuit<RANGE> {
+            type Config = Config;
+            type FloorPlanner = SimpleFloorPlanner;
+
+            fn without_witnesses(&self) -> Self {
+                MyCircuit(self.0)
+            }
+
+            fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
+                let selector = meta.selector();
+                let advice = meta.advice_column();
+
+                meta.create_gate("range check", |meta| {
+                    let selector = meta.query_selector(selector);
+                    let advice = meta.query_advice(advice, Rotation::cur());
+
+                    vec![selector * range_check(advice, RANGE)]
+                });
+
+                Config { selector, advice }
+            }
+
+            fn synthesize(
+                &self,
+                config: Self::Config,
+                mut layouter: impl Layouter<pallas::Base>,
+            ) -> Result<(), Error> {
+                layouter.assign_region(
+                    || "range constrain",
+                    |mut region| {
+                        config.selector.enable(&mut region, 0)?;
+                        region.assign_advice(
+                            || format!("witness {}", self.0),
+                            config.advice,
+                            0,
+                            || Ok(pallas::Base::from_u64(self.0.into())),
+                        )?;
+
+                        Ok(())
+                    },
+                )
+            }
+        }
+
+        for i in 0..8 {
+            let circuit: MyCircuit<8> = MyCircuit(i);
+            let prover = MockProver::<pallas::Base>::run(3, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        {
+            let circuit: MyCircuit<8> = MyCircuit(8);
+            let prover = MockProver::<pallas::Base>::run(3, &circuit, vec![]).unwrap();
+            assert_eq!(
+                prover.verify(),
+                Err(vec![VerifyFailure::ConstraintNotSatisfied {
+                    constraint: ((0, "range check").into(), 0, "").into(),
+                    row: 0
+                }])
+            );
+        }
+    }
+
+    #[test]
+    fn test_bitrange_subset() {
+        // Subset full range.
+        {
+            let field_elem = pallas::Base::rand();
+            let bitrange = 0..(pallas::Base::NUM_BITS as usize);
+            let subset = bitrange_subset(field_elem, bitrange);
+            assert_eq!(field_elem, subset);
+        }
+
+        // Subset zero bits
+        {
+            let field_elem = pallas::Base::rand();
+            let bitrange = 0..0;
+            let subset = bitrange_subset(field_elem, bitrange);
+            assert_eq!(pallas::Base::zero(), subset);
+        }
+
+        // Closure to decompose field element into pieces using consecutive ranges,
+        // and check that we recover the original.
+        let decompose = |field_elem: pallas::Base, ranges: &[Range<usize>]| {
+            assert_eq!(
+                ranges.iter().map(|range| range.len()).sum::<usize>(),
+                pallas::Base::NUM_BITS as usize
+            );
+            assert_eq!(ranges[0].start, 0);
+            assert_eq!(ranges.last().unwrap().end, pallas::Base::NUM_BITS as usize);
+
+            // Check ranges are contiguous
+            #[allow(unused_assignments)]
+            {
+                let mut ranges = ranges.iter();
+                let mut range = ranges.next().unwrap();
+                if let Some(next_range) = ranges.next() {
+                    assert_eq!(range.end, next_range.start);
+                    range = next_range;
+                }
+            }
+
+            let subsets = ranges
+                .iter()
+                .map(|range| bitrange_subset(field_elem, range.clone()))
+                .collect::<Vec<_>>();
+
+            let mut sum = subsets[0];
+            let mut num_bits = 0;
+            for (idx, subset) in subsets.iter().skip(1).enumerate() {
+                // 2^num_bits
+                let range_shift: [u8; 32] = {
+                    num_bits += ranges[idx].len();
+                    let mut range_shift = [0u8; 32];
+                    U256([2, 0, 0, 0])
+                        .pow(U256([num_bits as u64, 0, 0, 0]))
+                        .to_little_endian(&mut range_shift);
+                    range_shift
+                };
+                sum += subset * pallas::Base::from_bytes(&range_shift).unwrap();
+            }
+            assert_eq!(field_elem, sum);
+        };
+
+        decompose(pallas::Base::rand(), &[0..255]);
+        decompose(pallas::Base::rand(), &[0..1, 1..255]);
+        decompose(pallas::Base::rand(), &[0..254, 254..255]);
+        decompose(pallas::Base::rand(), &[0..127, 127..255]);
+        decompose(pallas::Base::rand(), &[0..128, 128..255]);
+        decompose(
+            pallas::Base::rand(),
+            &[0..50, 50..100, 100..150, 150..200, 200..255],
+        );
+    }
+}

examples/halo2/src/circuit/gadget/utilities/cond_swap.rs

@@ -0,0 +1,301 @@
+use super::{copy, CellValue, UtilitiesInstructions, Var};
+use halo2::{
+    circuit::{Chip, Layouter},
+    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector},
+    poly::Rotation,
+};
+use pasta_curves::arithmetic::FieldExt;
+use std::{array, marker::PhantomData};
+
+pub trait CondSwapInstructions<F: FieldExt>: UtilitiesInstructions<F> {
+    #[allow(clippy::type_complexity)]
+    /// Given an input pair (a,b) and a `swap` boolean flag, returns
+    /// (b,a) if `swap` is set, else (a,b) if `swap` is not set.
+    ///
+    /// The second element of the pair is required to be a witnessed
+    /// value, not a variable that already exists in the circuit.
+    fn swap(
+        &self,
+        layouter: impl Layouter<F>,
+        pair: (Self::Var, Option<F>),
+        swap: Option<bool>,
+    ) -> Result<(Self::Var, Self::Var), Error>;
+}
+
+/// A chip implementing a conditional swap.
+#[derive(Clone, Debug)]
+pub struct CondSwapChip<F> {
+    config: CondSwapConfig,
+    _marker: PhantomData<F>,
+}
+
+impl<F: FieldExt> Chip<F> for CondSwapChip<F> {
+    type Config = CondSwapConfig;
+    type Loaded = ();
+
+    fn config(&self) -> &Self::Config {
+        &self.config
+    }
+
+    fn loaded(&self) -> &Self::Loaded {
+        &()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct CondSwapConfig {
+    pub q_swap: Selector,
+    pub a: Column<Advice>,
+    pub b: Column<Advice>,
+    pub a_swapped: Column<Advice>,
+    pub b_swapped: Column<Advice>,
+    pub swap: Column<Advice>,
+}
+
+impl<F: FieldExt> UtilitiesInstructions<F> for CondSwapChip<F> {
+    type Var = CellValue<F>;
+}
+
+impl<F: FieldExt> CondSwapInstructions<F> for CondSwapChip<F> {
+    #[allow(clippy::type_complexity)]
+    fn swap(
+        &self,
+        mut layouter: impl Layouter<F>,
+        pair: (Self::Var, Option<F>),
+        swap: Option<bool>,
+    ) -> Result<(Self::Var, Self::Var), Error> {
+        let config = self.config();
+
+        layouter.assign_region(
+            || "swap",
+            |mut region| {
+                // Enable `q_swap` selector
+                config.q_swap.enable(&mut region, 0)?;
+
+                // Copy in `a` value
+                let a = copy(&mut region, || "copy a", config.a, 0, &pair.0)?;
+
+                // Witness `b` value
+                let b = {
+                    let cell = region.assign_advice(
+                        || "witness b",
+                        config.b,
+                        0,
+                        || pair.1.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue::new(cell, pair.1)
+                };
+
+                // Witness `swap` value
+                let swap_val = swap.map(|swap| F::from_u64(swap as u64));
+                region.assign_advice(
+                    || "swap",
+                    config.swap,
+                    0,
+                    || swap_val.ok_or(Error::SynthesisError),
+                )?;
+
+                // Conditionally swap a
+                let a_swapped = {
+                    let a_swapped = a
+                        .value
+                        .zip(b.value)
+                        .zip(swap)
+                        .map(|((a, b), swap)| if swap { b } else { a });
+                    let a_swapped_cell = region.assign_advice(
+                        || "a_swapped",
+                        config.a_swapped,
+                        0,
+                        || a_swapped.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue {
+                        cell: a_swapped_cell,
+                        value: a_swapped,
+                    }
+                };
+
+                // Conditionally swap b
+                let b_swapped = {
+                    let b_swapped = a
+                        .value
+                        .zip(b.value)
+                        .zip(swap)
+                        .map(|((a, b), swap)| if swap { a } else { b });
+                    let b_swapped_cell = region.assign_advice(
+                        || "b_swapped",
+                        config.b_swapped,
+                        0,
+                        || b_swapped.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue {
+                        cell: b_swapped_cell,
+                        value: b_swapped,
+                    }
+                };
+
+                // Return swapped pair
+                Ok((a_swapped, b_swapped))
+            },
+        )
+    }
+}
+
+impl<F: FieldExt> CondSwapChip<F> {
+    /// Configures this chip for use in a circuit.
+    ///
+    /// # Side-effects
+    ///
+    /// `advices[0]` will be equality-enabled.
+    pub fn configure(
+        meta: &mut ConstraintSystem<F>,
+        advices: [Column<Advice>; 5],
+    ) -> CondSwapConfig {
+        let a = advices[0];
+        // Only column a is used in an equality constraint directly by this chip.
+        meta.enable_equality(a.into());
+
+        let q_swap = meta.selector();
+
+        let config = CondSwapConfig {
+            q_swap,
+            a,
+            b: advices[1],
+            a_swapped: advices[2],
+            b_swapped: advices[3],
+            swap: advices[4],
+        };
+
+        // TODO: optimise shape of gate for Merkle path validation
+
+        meta.create_gate("a' = b ⋅ swap + a ⋅ (1-swap)", |meta| {
+            let q_swap = meta.query_selector(q_swap);
+
+            let a = meta.query_advice(config.a, Rotation::cur());
+            let b = meta.query_advice(config.b, Rotation::cur());
+            let a_swapped = meta.query_advice(config.a_swapped, Rotation::cur());
+            let b_swapped = meta.query_advice(config.b_swapped, Rotation::cur());
+            let swap = meta.query_advice(config.swap, Rotation::cur());
+
+            let one = Expression::Constant(F::one());
+
+            // a_swapped - b ⋅ swap - a ⋅ (1-swap) = 0
+            // This checks that `a_swapped` is equal to `y` when `swap` is set,
+            // but remains as `a` when `swap` is not set.
+            let a_check =
+                a_swapped - b.clone() * swap.clone() - a.clone() * (one.clone() - swap.clone());
+
+            // b_swapped - a ⋅ swap - b ⋅ (1-swap) = 0
+            // This checks that `b_swapped` is equal to `a` when `swap` is set,
+            // but remains as `b` when `swap` is not set.
+            let b_check = b_swapped - a * swap.clone() - b * (one.clone() - swap.clone());
+
+            // Check `swap` is boolean.
+            let bool_check = swap.clone() * (one - swap);
+
+            array::IntoIter::new([a_check, b_check, bool_check])
+                .map(move |poly| q_swap.clone() * poly)
+        });
+
+        config
+    }
+
+    pub fn construct(config: CondSwapConfig) -> Self {
+        CondSwapChip {
+            config,
+            _marker: PhantomData,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::UtilitiesInstructions;
+    use super::{CondSwapChip, CondSwapConfig, CondSwapInstructions};
+    use halo2::{
+        circuit::{Layouter, SimpleFloorPlanner},
+        dev::MockProver,
+        plonk::{Circuit, ConstraintSystem, Error},
+    };
+    use pasta_curves::{arithmetic::FieldExt, pallas::Base};
+
+    #[test]
+    fn cond_swap() {
+        #[derive(Default)]
+        struct MyCircuit<F: FieldExt> {
+            a: Option<F>,
+            b: Option<F>,
+            swap: Option<bool>,
+        }
+
+        impl<F: FieldExt> Circuit<F> for MyCircuit<F> {
+            type Config = CondSwapConfig;
+            type FloorPlanner = SimpleFloorPlanner;
+
+            fn without_witnesses(&self) -> Self {
+                Self::default()
+            }
+
+            fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
+                let advices = [
+                    meta.advice_column(),
+                    meta.advice_column(),
+                    meta.advice_column(),
+                    meta.advice_column(),
+                    meta.advice_column(),
+                ];
+
+                CondSwapChip::<F>::configure(meta, advices)
+            }
+
+            fn synthesize(
+                &self,
+                config: Self::Config,
+                mut layouter: impl Layouter<F>,
+            ) -> Result<(), Error> {
+                let chip = CondSwapChip::<F>::construct(config.clone());
+
+                // Load the pair and the swap flag into the circuit.
+                let a = chip.load_private(layouter.namespace(|| "a"), config.a, self.a)?;
+                // Return the swapped pair.
+                let swapped_pair =
+                    chip.swap(layouter.namespace(|| "swap"), (a, self.b), self.swap)?;
+
+                if let Some(swap) = self.swap {
+                    if swap {
+                        // Check that `a` and `b` have been swapped
+                        assert_eq!(swapped_pair.0.value.unwrap(), self.b.unwrap());
+                        assert_eq!(swapped_pair.1.value.unwrap(), a.value.unwrap());
+                    } else {
+                        // Check that `a` and `b` have not been swapped
+                        assert_eq!(swapped_pair.0.value.unwrap(), a.value.unwrap());
+                        assert_eq!(swapped_pair.1.value.unwrap(), self.b.unwrap());
+                    }
+                }
+
+                Ok(())
+            }
+        }
+
+        // Test swap case
+        {
+            let circuit: MyCircuit<Base> = MyCircuit {
+                a: Some(Base::rand()),
+                b: Some(Base::rand()),
+                swap: Some(true),
+            };
+            let prover = MockProver::<Base>::run(3, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // Test non-swap case
+        {
+            let circuit: MyCircuit<Base> = MyCircuit {
+                a: Some(Base::rand()),
+                b: Some(Base::rand()),
+                swap: Some(false),
+            };
+            let prover = MockProver::<Base>::run(3, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+    }
+}

examples/halo2/src/circuit/gadget/utilities/decompose_running_sum.rs

@@ -0,0 +1,393 @@
+//! Decomposes an $n$-bit field element $\alpha$ into $W$ windows, each window
+//! being a $K$-bit word, using a running sum $z$.
+//! We constrain $K \leq 3$ for this helper.
+//!     $$\alpha = k_0 + (2^K) k_1 + (2^{2K}) k_2 + ... + (2^{(W-1)K}) k_{W-1}$$
+//!
+//! $z_0$ is initialized as $\alpha$. Each successive $z_{i+1}$ is computed as
+//!                $$z_{i+1} = (z_{i} - k_i) / (2^K).$$
+//! $z_W$ is constrained to be zero.
+//! The difference between each interstitial running sum output is constrained
+//! to be $K$ bits, i.e.
+//!                      `range_check`($k_i$, $2^K$),
+//! where
+//! ```text
+//!   range_check(word, range)
+//!     = word * (1 - word) * (2 - word) * ... * ((range - 1) - word)
+//! ```
+//!
+//! Given that the `range_check` constraint will be toggled by a selector, in
+//! practice we will have a `selector * range_check(word, range)` expression
+//! of degree `range + 1`.
+//!
+//! This means that $2^K$ has to be at most `degree_bound - 1` in order for
+//! the range check constraint to stay within the degree bound.
+
+use ff::PrimeFieldBits;
+use halo2::{
+    circuit::Region,
+    plonk::{Advice, Column, ConstraintSystem, Error, Selector},
+    poly::Rotation,
+};
+
+use super::{copy, range_check, CellValue, Var};
+use crate::constants::util::decompose_word;
+use pasta_curves::arithmetic::FieldExt;
+use std::marker::PhantomData;
+
+/// The running sum $[z_0, ..., z_W]$. If created in strict mode, $z_W = 0$.
+pub struct RunningSum<F: FieldExt + PrimeFieldBits>(Vec<CellValue<F>>);
+impl<F: FieldExt + PrimeFieldBits> std::ops::Deref for RunningSum<F> {
+    type Target = Vec<CellValue<F>>;
+
+    fn deref(&self) -> &Vec<CellValue<F>> {
+        &self.0
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct RunningSumConfig<F: FieldExt + PrimeFieldBits, const WINDOW_NUM_BITS: usize> {
+    q_range_check: Selector,
+    pub z: Column<Advice>,
+    _marker: PhantomData<F>,
+}
+
+impl<F: FieldExt + PrimeFieldBits, const WINDOW_NUM_BITS: usize>
+    RunningSumConfig<F, WINDOW_NUM_BITS>
+{
+    /// `perm` MUST include the advice column `z`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if WINDOW_NUM_BITS > 3.
+    ///
+    /// # Side-effects
+    ///
+    /// `z` will be equality-enabled.
+    pub fn configure(
+        meta: &mut ConstraintSystem<F>,
+        q_range_check: Selector,
+        z: Column<Advice>,
+    ) -> Self {
+        assert!(WINDOW_NUM_BITS <= 3);
+
+        meta.enable_equality(z.into());
+
+        let config = Self {
+            q_range_check,
+            z,
+            _marker: PhantomData,
+        };
+
+        meta.create_gate("range check", |meta| {
+            let q_range_check = meta.query_selector(config.q_range_check);
+            let z_cur = meta.query_advice(config.z, Rotation::cur());
+            let z_next = meta.query_advice(config.z, Rotation::next());
+            //    z_i = 2^{K}⋅z_{i + 1} + k_i
+            // => k_i = z_i - 2^{K}⋅z_{i + 1}
+            let word = z_cur - z_next * F::from_u64(1 << WINDOW_NUM_BITS);
+
+            vec![q_range_check * range_check(word, 1 << WINDOW_NUM_BITS)]
+        });
+
+        config
+    }
+
+    /// Decompose a field element alpha that is witnessed in this helper.
+    ///
+    /// `strict` = true constrains the final running sum to be zero, i.e.
+    /// constrains alpha to be within WINDOW_NUM_BITS * num_windows bits.
+    pub fn witness_decompose(
+        &self,
+        region: &mut Region<'_, F>,
+        offset: usize,
+        alpha: Option<F>,
+        strict: bool,
+        word_num_bits: usize,
+        num_windows: usize,
+    ) -> Result<RunningSum<F>, Error> {
+        let z_0 = {
+            let cell = region.assign_advice(
+                || "z_0 = alpha",
+                self.z,
+                offset,
+                || alpha.ok_or(Error::SynthesisError),
+            )?;
+            CellValue::new(cell, alpha)
+        };
+        self.decompose(region, offset, z_0, strict, word_num_bits, num_windows)
+    }
+
+    /// Decompose an existing variable alpha that is copied into this helper.
+    ///
+    /// `strict` = true constrains the final running sum to be zero, i.e.
+    /// constrains alpha to be within WINDOW_NUM_BITS * num_windows bits.
+    pub fn copy_decompose(
+        &self,
+        region: &mut Region<'_, F>,
+        offset: usize,
+        alpha: CellValue<F>,
+        strict: bool,
+        word_num_bits: usize,
+        num_windows: usize,
+    ) -> Result<RunningSum<F>, Error> {
+        let z_0 = copy(region, || "copy z_0 = alpha", self.z, offset, &alpha)?;
+        self.decompose(region, offset, z_0, strict, word_num_bits, num_windows)
+    }
+
+    /// `z_0` must be the cell at `(self.z, offset)` in `region`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if there are too many windows for the given word size.
+    fn decompose(
+        &self,
+        region: &mut Region<'_, F>,
+        offset: usize,
+        z_0: CellValue<F>,
+        strict: bool,
+        word_num_bits: usize,
+        num_windows: usize,
+    ) -> Result<RunningSum<F>, Error> {
+        // Make sure that we do not have more windows than required for the number
+        // of bits in the word. In other words, every window must contain at least
+        // one bit of the word (no empty windows).
+        //
+        // For example, let:
+        //      - word_num_bits = 64
+        //      - WINDOW_NUM_BITS = 3
+        // In this case, the maximum allowed num_windows is 22:
+        //                    3 * 22 < 64 + 3
+        //
+        assert!(WINDOW_NUM_BITS * num_windows < word_num_bits + WINDOW_NUM_BITS);
+
+        // Enable selectors
+        for idx in 0..num_windows {
+            self.q_range_check.enable(region, offset + idx)?;
+        }
+
+        // Decompose base field element into K-bit words.
+        let words: Vec<Option<u8>> = {
+            let words = z_0
+                .value()
+                .map(|word| decompose_word::<F>(word, word_num_bits, WINDOW_NUM_BITS));
+
+            if let Some(words) = words {
+                words.into_iter().map(Some).collect()
+            } else {
+                vec![None; num_windows]
+            }
+        };
+
+        // Initialize empty vector to store running sum values [z_0, ..., z_W].
+        let mut zs: Vec<CellValue<F>> = vec![z_0];
+        let mut z = z_0;
+
+        // Assign running sum `z_{i+1}` = (z_i - k_i) / (2^K) for i = 0..=n-1.
+        // Outside of this helper, z_0 = alpha must have already been loaded into the
+        // `z` column at `offset`.
+        let two_pow_k_inv = F::from_u64(1 << WINDOW_NUM_BITS as u64).invert().unwrap();
+        for (i, word) in words.iter().enumerate() {
+            // z_next = (z_cur - word) / (2^K)
+            let z_next = {
+                let word = word.map(|word| F::from_u64(word as u64));
+                let z_next_val = z
+                    .value()
+                    .zip(word)
+                    .map(|(z_cur_val, word)| (z_cur_val - word) * two_pow_k_inv);
+                let cell = region.assign_advice(
+                    || format!("z_{:?}", i + 1),
+                    self.z,
+                    offset + i + 1,
+                    || z_next_val.ok_or(Error::SynthesisError),
+                )?;
+                CellValue::new(cell, z_next_val)
+            };
+
+            // Update `z`.
+            z = z_next;
+            zs.push(z);
+        }
+        assert_eq!(zs.len(), num_windows + 1);
+
+        if strict {
+            // Constrain the final running sum output to be zero.
+            region.constrain_constant(zs.last().unwrap().cell(), F::zero())?;
+        }
+
+        Ok(RunningSum(zs))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::constants::{self, FIXED_BASE_WINDOW_SIZE, L_ORCHARD_BASE, L_VALUE};
+    use halo2::{
+        circuit::{Layouter, SimpleFloorPlanner},
+        dev::{MockProver, VerifyFailure},
+        plonk::{Any, Circuit, ConstraintSystem, Error},
+    };
+    use pasta_curves::{arithmetic::FieldExt, pallas};
+
+    #[test]
+    fn test_running_sum() {
+        struct MyCircuit<
+            F: FieldExt + PrimeFieldBits,
+            const WORD_NUM_BITS: usize,
+            const WINDOW_NUM_BITS: usize,
+            const NUM_WINDOWS: usize,
+        > {
+            alpha: Option<F>,
+            strict: bool,
+        }
+
+        impl<
+                F: FieldExt + PrimeFieldBits,
+                const WORD_NUM_BITS: usize,
+                const WINDOW_NUM_BITS: usize,
+                const NUM_WINDOWS: usize,
+            > Circuit<F> for MyCircuit<F, WORD_NUM_BITS, WINDOW_NUM_BITS, NUM_WINDOWS>
+        {
+            type Config = RunningSumConfig<F, WINDOW_NUM_BITS>;
+            type FloorPlanner = SimpleFloorPlanner;
+
+            fn without_witnesses(&self) -> Self {
+                Self {
+                    alpha: None,
+                    strict: self.strict,
+                }
+            }
+
+            fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
+                let z = meta.advice_column();
+                let q_range_check = meta.selector();
+                let constants = meta.fixed_column();
+                meta.enable_constant(constants);
+
+                RunningSumConfig::<F, WINDOW_NUM_BITS>::configure(meta, q_range_check, z)
+            }
+
+            fn synthesize(
+                &self,
+                config: Self::Config,
+                mut layouter: impl Layouter<F>,
+            ) -> Result<(), Error> {
+                layouter.assign_region(
+                    || "decompose",
+                    |mut region| {
+                        let offset = 0;
+                        let zs = config.witness_decompose(
+                            &mut region,
+                            offset,
+                            self.alpha,
+                            self.strict,
+                            WORD_NUM_BITS,
+                            NUM_WINDOWS,
+                        )?;
+                        let alpha = zs[0];
+
+                        let offset = offset + NUM_WINDOWS + 1;
+
+                        config.copy_decompose(
+                            &mut region,
+                            offset,
+                            alpha,
+                            self.strict,
+                            WORD_NUM_BITS,
+                            NUM_WINDOWS,
+                        )?;
+
+                        Ok(())
+                    },
+                )
+            }
+        }
+
+        // Random base field element
+        {
+            let alpha = pallas::Base::rand();
+
+            // Strict full decomposition should pass.
+            let circuit: MyCircuit<
+                pallas::Base,
+                L_ORCHARD_BASE,
+                FIXED_BASE_WINDOW_SIZE,
+                { constants::NUM_WINDOWS },
+            > = MyCircuit {
+                alpha: Some(alpha),
+                strict: true,
+            };
+            let prover = MockProver::<pallas::Base>::run(8, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // Random 64-bit word
+        {
+            let alpha = pallas::Base::from_u64(rand::random());
+
+            // Strict full decomposition should pass.
+            let circuit: MyCircuit<
+                pallas::Base,
+                L_VALUE,
+                FIXED_BASE_WINDOW_SIZE,
+                { constants::NUM_WINDOWS_SHORT },
+            > = MyCircuit {
+                alpha: Some(alpha),
+                strict: true,
+            };
+            let prover = MockProver::<pallas::Base>::run(8, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // 2^66
+        {
+            let alpha = pallas::Base::from_u128(1 << 66);
+
+            // Strict partial decomposition should fail.
+            let circuit: MyCircuit<
+                pallas::Base,
+                L_VALUE,
+                FIXED_BASE_WINDOW_SIZE,
+                { constants::NUM_WINDOWS_SHORT },
+            > = MyCircuit {
+                alpha: Some(alpha),
+                strict: true,
+            };
+            let prover = MockProver::<pallas::Base>::run(8, &circuit, vec![]).unwrap();
+            assert_eq!(
+                prover.verify(),
+                Err(vec![
+                    VerifyFailure::Permutation {
+                        column: (Any::Fixed, 0).into(),
+                        row: 0
+                    },
+                    VerifyFailure::Permutation {
+                        column: (Any::Fixed, 0).into(),
+                        row: 1
+                    },
+                    VerifyFailure::Permutation {
+                        column: (Any::Advice, 0).into(),
+                        row: 22
+                    },
+                    VerifyFailure::Permutation {
+                        column: (Any::Advice, 0).into(),
+                        row: 45
+                    },
+                ])
+            );
+
+            // Non-strict partial decomposition should pass.
+            let circuit: MyCircuit<
+                pallas::Base,
+                { constants::L_VALUE },
+                FIXED_BASE_WINDOW_SIZE,
+                { constants::NUM_WINDOWS_SHORT },
+            > = MyCircuit {
+                alpha: Some(alpha),
+                strict: false,
+            };
+            let prover = MockProver::<pallas::Base>::run(8, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+    }
+}

examples/halo2/src/circuit/gadget/utilities/lookup_range_check.rs

@@ -0,0 +1,629 @@
+//! Make use of a K-bit lookup table to decompose a field element into K-bit
+//! words.
+
+use crate::spec::lebs2ip;
+use halo2::{
+    circuit::{Layouter, Region},
+    plonk::{Advice, Column, ConstraintSystem, Error, Selector, TableColumn},
+    poly::Rotation,
+};
+use std::{convert::TryInto, marker::PhantomData};
+
+use ff::PrimeFieldBits;
+
+use super::*;
+
+/// The running sum $[z_0, ..., z_W]$. If created in strict mode, $z_W = 0$.
+pub struct RunningSum<F: FieldExt + PrimeFieldBits>(Vec<CellValue<F>>);
+impl<F: FieldExt + PrimeFieldBits> std::ops::Deref for RunningSum<F> {
+    type Target = Vec<CellValue<F>>;
+
+    fn deref(&self) -> &Vec<CellValue<F>> {
+        &self.0
+    }
+}
+
+#[derive(Eq, PartialEq, Debug, Clone)]
+pub struct LookupRangeCheckConfig<F: FieldExt + PrimeFieldBits, const K: usize> {
+    pub q_lookup: Selector,
+    pub q_running: Selector,
+    pub q_bitshift: Selector,
+    pub running_sum: Column<Advice>,
+    table_idx: TableColumn,
+    _marker: PhantomData<F>,
+}
+
+impl<F: FieldExt + PrimeFieldBits, const K: usize> LookupRangeCheckConfig<F, K> {
+    /// The `running_sum` advice column breaks the field element into `K`-bit
+    /// words. It is used to construct the input expression to the lookup
+    /// argument.
+    ///
+    /// The `table_idx` fixed column contains values from [0..2^K). Looking up
+    /// a value in `table_idx` constrains it to be within this range. The table
+    /// can be loaded outside this helper.
+    ///
+    /// # Side-effects
+    ///
+    /// Both the `running_sum` and `constants` columns will be equality-enabled.
+    pub fn configure(
+        meta: &mut ConstraintSystem<F>,
+        running_sum: Column<Advice>,
+        table_idx: TableColumn,
+    ) -> Self {
+        meta.enable_equality(running_sum.into());
+
+        let q_lookup = meta.complex_selector();
+        let q_running = meta.complex_selector();
+        let q_bitshift = meta.selector();
+        let config = LookupRangeCheckConfig {
+            q_lookup,
+            q_running,
+            q_bitshift,
+            running_sum,
+            table_idx,
+            _marker: PhantomData,
+        };
+
+        meta.lookup(|meta| {
+            let q_lookup = meta.query_selector(config.q_lookup);
+            let q_running = meta.query_selector(config.q_running);
+            let z_cur = meta.query_advice(config.running_sum, Rotation::cur());
+
+            // In the case of a running sum decomposition, we recover the word from
+            // the difference of the running sums:
+            //    z_i = 2^{K}⋅z_{i + 1} + a_i
+            // => a_i = z_i - 2^{K}⋅z_{i + 1}
+            let running_sum_lookup = {
+                let running_sum_word = {
+                    let z_next = meta.query_advice(config.running_sum, Rotation::next());
+                    z_cur.clone() - z_next * F::from_u64(1 << K)
+                };
+
+                q_running.clone() * running_sum_word
+            };
+
+            // In the short range check, the word is directly witnessed.
+            let short_lookup = {
+                let short_word = z_cur;
+                let q_short = Expression::Constant(F::one()) - q_running;
+
+                q_short * short_word
+            };
+
+            // Combine the running sum and short lookups:
+            vec![(
+                q_lookup * (running_sum_lookup + short_lookup),
+                config.table_idx,
+            )]
+        });
+
+        // For short lookups, check that the word has been shifted by the correct number of bits.
+        meta.create_gate("Short lookup bitshift", |meta| {
+            let q_bitshift = meta.query_selector(config.q_bitshift);
+            let word = meta.query_advice(config.running_sum, Rotation::prev());
+            let shifted_word = meta.query_advice(config.running_sum, Rotation::cur());
+            let inv_two_pow_s = meta.query_advice(config.running_sum, Rotation::next());
+
+            let two_pow_k = F::from_u64(1 << K);
+
+            // shifted_word = word * 2^{K-s}
+            //              = word * 2^K * inv_two_pow_s
+            vec![q_bitshift * (word * two_pow_k * inv_two_pow_s - shifted_word)]
+        });
+
+        config
+    }
+
+    #[cfg(test)]
+    // Loads the values [0..2^K) into `table_idx`. This is only used in testing
+    // for now, since the Sinsemilla chip provides a pre-loaded table in the
+    // Orchard context.
+    pub fn load(&self, layouter: &mut impl Layouter<F>) -> Result<(), Error> {
+        layouter.assign_table(
+            || "table_idx",
+            |mut table| {
+                // We generate the row values lazily (we only need them during keygen).
+                for index in 0..(1 << K) {
+                    table.assign_cell(
+                        || "table_idx",
+                        self.table_idx,
+                        index,
+                        || Ok(F::from_u64(index as u64)),
+                    )?;
+                }
+                Ok(())
+            },
+        )
+    }
+
+    /// Range check on an existing cell that is copied into this helper.
+    ///
+    /// Returns an error if `element` is not in a column that was passed to
+    /// [`ConstraintSystem::enable_equality`] during circuit configuration.
+    pub fn copy_check(
+        &self,
+        mut layouter: impl Layouter<F>,
+        element: CellValue<F>,
+        num_words: usize,
+        strict: bool,
+    ) -> Result<RunningSum<F>, Error> {
+        layouter.assign_region(
+            || format!("{:?} words range check", num_words),
+            |mut region| {
+                // Copy `element` and initialize running sum `z_0 = element` to decompose it.
+                let z_0 = copy(&mut region, || "z_0", self.running_sum, 0, &element)?;
+                self.range_check(&mut region, z_0, num_words, strict)
+            },
+        )
+    }
+
+    /// Range check on a value that is witnessed in this helper.
+    pub fn witness_check(
+        &self,
+        mut layouter: impl Layouter<F>,
+        value: Option<F>,
+        num_words: usize,
+        strict: bool,
+    ) -> Result<RunningSum<F>, Error> {
+        layouter.assign_region(
+            || "Witness element",
+            |mut region| {
+                let z_0 = {
+                    let cell = region.assign_advice(
+                        || "Witness element",
+                        self.running_sum,
+                        0,
+                        || value.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue::new(cell, value)
+                };
+                self.range_check(&mut region, z_0, num_words, strict)
+            },
+        )
+    }
+
+    /// If `strict` is set to "true", the field element must fit into
+    /// `num_words * K` bits. In other words, the the final cumulative sum `z_{num_words}`
+    /// must be zero.
+    ///
+    /// If `strict` is set to "false", the final `z_{num_words}` is not constrained.
+    ///
+    /// `element` must have been assigned to `self.running_sum` at offset 0.
+    fn range_check(
+        &self,
+        region: &mut Region<'_, F>,
+        element: CellValue<F>,
+        num_words: usize,
+        strict: bool,
+    ) -> Result<RunningSum<F>, Error> {
+        // `num_words` must fit into a single field element.
+        assert!(num_words * K <= F::CAPACITY as usize);
+        let num_bits = num_words * K;
+
+        // Chunk the first num_bits bits into K-bit words.
+        let words = {
+            // Take first num_bits bits of `element`.
+            let bits = element.value().map(|element| {
+                element
+                    .to_le_bits()
+                    .into_iter()
+                    .take(num_bits)
+                    .collect::<Vec<_>>()
+            });
+
+            let words: Option<Vec<F>> = bits.map(|bits| {
+                bits.chunks_exact(K)
+                    .map(|word| F::from_u64(lebs2ip::<K>(&(word.try_into().unwrap()))))
+                    .collect::<Vec<_>>()
+            });
+            if let Some(words) = words {
+                words.into_iter().map(Some).collect()
+            } else {
+                vec![None; num_words]
+            }
+        };
+
+        let mut zs = vec![element];
+
+        // Assign cumulative sum such that
+        //          z_i = 2^{K}⋅z_{i + 1} + a_i
+        // => z_{i + 1} = (z_i - a_i) / (2^K)
+        //
+        // For `element` = a_0 + 2^10 a_1 + ... + 2^{120} a_{12}}, initialize z_0 = `element`.
+        // If `element` fits in 130 bits, we end up with z_{13} = 0.
+        let mut z = element;
+        let inv_two_pow_k = F::from_u64(1u64 << K).invert().unwrap();
+        for (idx, word) in words.iter().enumerate() {
+            // Enable q_lookup on this row
+            self.q_lookup.enable(region, idx)?;
+            // Enable q_running on this row
+            self.q_running.enable(region, idx)?;
+
+            // z_next = (z_cur - m_cur) / 2^K
+            z = {
+                let z_val = z
+                    .value()
+                    .zip(*word)
+                    .map(|(z, word)| (z - word) * inv_two_pow_k);
+
+                // Assign z_next
+                let z_cell = region.assign_advice(
+                    || format!("z_{:?}", idx + 1),
+                    self.running_sum,
+                    idx + 1,
+                    || z_val.ok_or(Error::SynthesisError),
+                )?;
+
+                CellValue::new(z_cell, z_val)
+            };
+            zs.push(z);
+        }
+
+        if strict {
+            // Constrain the final `z` to be zero.
+            region.constrain_constant(zs.last().unwrap().cell(), F::zero())?;
+        }
+
+        Ok(RunningSum(zs))
+    }
+
+    /// Short range check on an existing cell that is copied into this helper.
+    ///
+    /// # Panics
+    ///
+    /// Panics if NUM_BITS is equal to or larger than K.
+    pub fn copy_short_check(
+        &self,
+        mut layouter: impl Layouter<F>,
+        element: CellValue<F>,
+        num_bits: usize,
+    ) -> Result<(), Error> {
+        assert!(num_bits < K);
+        layouter.assign_region(
+            || format!("Range check {:?} bits", num_bits),
+            |mut region| {
+                // Copy `element` to use in the k-bit lookup.
+                let element = copy(&mut region, || "element", self.running_sum, 0, &element)?;
+
+                self.short_range_check(&mut region, element, num_bits)
+            },
+        )
+    }
+
+    /// Short range check on value that is witnessed in this helper.
+    ///
+    /// # Panics
+    ///
+    /// Panics if num_bits is larger than K.
+    pub fn witness_short_check(
+        &self,
+        mut layouter: impl Layouter<F>,
+        element: Option<F>,
+        num_bits: usize,
+    ) -> Result<CellValue<F>, Error> {
+        assert!(num_bits <= K);
+        layouter.assign_region(
+            || format!("Range check {:?} bits", num_bits),
+            |mut region| {
+                // Witness `element` to use in the k-bit lookup.
+                let element = {
+                    let cell = region.assign_advice(
+                        || "Witness element",
+                        self.running_sum,
+                        0,
+                        || element.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue::new(cell, element)
+                };
+
+                self.short_range_check(&mut region, element, num_bits)?;
+
+                Ok(element)
+            },
+        )
+    }
+
+    /// Constrain `x` to be a NUM_BITS word.
+    ///
+    /// `element` must have been assigned to `self.running_sum` at offset 0.
+    fn short_range_check(
+        &self,
+        region: &mut Region<'_, F>,
+        element: CellValue<F>,
+        num_bits: usize,
+    ) -> Result<(), Error> {
+        // Enable lookup for `element`, to constrain it to 10 bits.
+        self.q_lookup.enable(region, 0)?;
+
+        // Enable lookup for shifted element, to constrain it to 10 bits.
+        self.q_lookup.enable(region, 1)?;
+
+        // Check element has been shifted by the correct number of bits.
+        self.q_bitshift.enable(region, 1)?;
+
+        // Assign shifted `element * 2^{K - num_bits}`
+        let shifted = element.value().map(|element| {
+            let shift = F::from_u64(1 << (K - num_bits));
+            element * shift
+        });
+
+        region.assign_advice(
+            || format!("element * 2^({}-{})", K, num_bits),
+            self.running_sum,
+            1,
+            || shifted.ok_or(Error::SynthesisError),
+        )?;
+
+        // Assign 2^{-num_bits} from a fixed column.
+        let inv_two_pow_s = F::from_u64(1 << num_bits).invert().unwrap();
+        region.assign_advice_from_constant(
+            || format!("2^(-{})", num_bits),
+            self.running_sum,
+            2,
+            inv_two_pow_s,
+        )?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::Var;
+    use super::LookupRangeCheckConfig;
+
+    use crate::primitives::sinsemilla::{INV_TWO_POW_K, K};
+    use crate::spec::lebs2ip;
+    use ff::{Field, PrimeFieldBits};
+    use halo2::{
+        circuit::{Layouter, SimpleFloorPlanner},
+        dev::{MockProver, VerifyFailure},
+        plonk::{Circuit, ConstraintSystem, Error},
+    };
+    use pasta_curves::{arithmetic::FieldExt, pallas};
+
+    use std::{convert::TryInto, marker::PhantomData};
+
+    #[test]
+    fn lookup_range_check() {
+        #[derive(Clone, Copy)]
+        struct MyCircuit<F: FieldExt + PrimeFieldBits> {
+            num_words: usize,
+            _marker: PhantomData<F>,
+        }
+
+        impl<F: FieldExt + PrimeFieldBits> Circuit<F> for MyCircuit<F> {
+            type Config = LookupRangeCheckConfig<F, K>;
+            type FloorPlanner = SimpleFloorPlanner;
+
+            fn without_witnesses(&self) -> Self {
+                *self
+            }
+
+            fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
+                let running_sum = meta.advice_column();
+                let table_idx = meta.lookup_table_column();
+                let constants = meta.fixed_column();
+                meta.enable_constant(constants);
+
+                LookupRangeCheckConfig::<F, K>::configure(meta, running_sum, table_idx)
+            }
+
+            fn synthesize(
+                &self,
+                config: Self::Config,
+                mut layouter: impl Layouter<F>,
+            ) -> Result<(), Error> {
+                // Load table_idx
+                config.load(&mut layouter)?;
+
+                // Lookup constraining element to be no longer than num_words * K bits.
+                let elements_and_expected_final_zs = [
+                    (
+                        F::from_u64((1 << (self.num_words * K)) - 1),
+                        F::zero(),
+                        true,
+                    ), // a word that is within self.num_words * K bits long
+                    (F::from_u64(1 << (self.num_words * K)), F::one(), false), // a word that is just over self.num_words * K bits long
+                ];
+
+                fn expected_zs<F: FieldExt + PrimeFieldBits, const K: usize>(
+                    element: F,
+                    num_words: usize,
+                ) -> Vec<F> {
+                    let chunks = {
+                        element
+                            .to_le_bits()
+                            .iter()
+                            .by_val()
+                            .take(num_words * K)
+                            .collect::<Vec<_>>()
+                            .chunks_exact(K)
+                            .map(|chunk| F::from_u64(lebs2ip::<K>(chunk.try_into().unwrap())))
+                            .collect::<Vec<_>>()
+                    };
+                    let expected_zs = {
+                        let inv_two_pow_k = F::from_bytes(&INV_TWO_POW_K).unwrap();
+                        chunks.iter().fold(vec![element], |mut zs, a_i| {
+                            // z_{i + 1} = (z_i - a_i) / 2^{K}
+                            let z = (zs[zs.len() - 1] - a_i) * inv_two_pow_k;
+                            zs.push(z);
+                            zs
+                        })
+                    };
+                    expected_zs
+                }
+
+                for (element, expected_final_z, strict) in elements_and_expected_final_zs.iter() {
+                    let expected_zs = expected_zs::<F, K>(*element, self.num_words);
+
+                    let zs = config.witness_check(
+                        layouter.namespace(|| format!("Lookup {:?}", self.num_words)),
+                        Some(*element),
+                        self.num_words,
+                        *strict,
+                    )?;
+
+                    assert_eq!(*expected_zs.last().unwrap(), *expected_final_z);
+
+                    for (expected_z, z) in expected_zs.into_iter().zip(zs.iter()) {
+                        if let Some(z) = z.value() {
+                            assert_eq!(expected_z, z);
+                        }
+                    }
+                }
+                Ok(())
+            }
+        }
+
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                num_words: 6,
+                _marker: PhantomData,
+            };
+
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+    }
+
+    #[test]
+    fn short_range_check() {
+        struct MyCircuit<F: FieldExt + PrimeFieldBits> {
+            element: Option<F>,
+            num_bits: usize,
+        }
+
+        impl<F: FieldExt + PrimeFieldBits> Circuit<F> for MyCircuit<F> {
+            type Config = LookupRangeCheckConfig<F, K>;
+            type FloorPlanner = SimpleFloorPlanner;
+
+            fn without_witnesses(&self) -> Self {
+                MyCircuit {
+                    element: None,
+                    num_bits: self.num_bits,
+                }
+            }
+
+            fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
+                let running_sum = meta.advice_column();
+                let table_idx = meta.lookup_table_column();
+                let constants = meta.fixed_column();
+                meta.enable_constant(constants);
+
+                LookupRangeCheckConfig::<F, K>::configure(meta, running_sum, table_idx)
+            }
+
+            fn synthesize(
+                &self,
+                config: Self::Config,
+                mut layouter: impl Layouter<F>,
+            ) -> Result<(), Error> {
+                // Load table_idx
+                config.load(&mut layouter)?;
+
+                // Lookup constraining element to be no longer than num_bits.
+                config.witness_short_check(
+                    layouter.namespace(|| format!("Lookup {:?} bits", self.num_bits)),
+                    self.element,
+                    self.num_bits,
+                )?;
+
+                Ok(())
+            }
+        }
+
+        // Edge case: zero bits
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(pallas::Base::zero()),
+                num_bits: 0,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // Edge case: K bits
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(pallas::Base::from_u64((1 << K) - 1)),
+                num_bits: K,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // Element within `num_bits`
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(pallas::Base::from_u64((1 << 6) - 1)),
+                num_bits: 6,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(prover.verify(), Ok(()));
+        }
+
+        // Element larger than `num_bits` but within K bits
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(pallas::Base::from_u64(1 << 6)),
+                num_bits: 6,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(
+                prover.verify(),
+                Err(vec![VerifyFailure::Lookup {
+                    lookup_index: 0,
+                    row: 1
+                }])
+            );
+        }
+
+        // Element larger than K bits
+        {
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(pallas::Base::from_u64(1 << K)),
+                num_bits: 6,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(
+                prover.verify(),
+                Err(vec![
+                    VerifyFailure::Lookup {
+                        lookup_index: 0,
+                        row: 0
+                    },
+                    VerifyFailure::Lookup {
+                        lookup_index: 0,
+                        row: 1
+                    },
+                ])
+            );
+        }
+
+        // Element which is not within `num_bits`, but which has a shifted value within
+        // num_bits
+        {
+            let num_bits = 6;
+            let shifted = pallas::Base::from_u64((1 << num_bits) - 1);
+            // Recall that shifted = element * 2^{K-s}
+            //          => element = shifted * 2^{s-K}
+            let element = shifted
+                * pallas::Base::from_u64(1 << (K as u64 - num_bits))
+                    .invert()
+                    .unwrap();
+            let circuit: MyCircuit<pallas::Base> = MyCircuit {
+                element: Some(element),
+                num_bits: num_bits as usize,
+            };
+            let prover = MockProver::<pallas::Base>::run(11, &circuit, vec![]).unwrap();
+            assert_eq!(
+                prover.verify(),
+                Err(vec![VerifyFailure::Lookup {
+                    lookup_index: 0,
+                    row: 0
+                }])
+            );
+        }
+    }
+}