1use crate::Vec;
  2use ark_ff::Field;
  3use ark_relations::r1cs::{Namespace, SynthesisError};
  4use core::borrow::Borrow;
  5
  6/// Describes the mode that a variable should be allocated in within
  7/// a `ConstraintSystem`.
  8#[derive(Eq, PartialEq, Ord, PartialOrd, Debug, Copy, Clone)]
  9pub enum AllocationMode {
 10    /// Indicate to the `ConstraintSystem` that the high-level variable should
 11    /// be allocated as a constant. That is, no `Variable`s should be
 12    /// generated.
 13    Constant = 0,
 14
 15    /// Indicate to the `ConstraintSystem` that the high-level variable should
 16    /// be allocated as a public input to the `ConstraintSystem`.
 17    Input = 1,
 18
 19    /// Indicate to the `ConstraintSystem` that the high-level variable should
 20    /// be allocated as a private witness to the `ConstraintSystem`.
 21    Witness = 2,
 22}
 23
 24impl AllocationMode {
 25    /// Outputs the maximum according to the relation `Constant < Input <
 26    /// Witness`.
 27    pub fn max(&self, other: Self) -> Self {
 28        use AllocationMode::*;
 29        match (self, other) {
 30            (Constant, _) => other,
 31            (Input, Constant) => *self,
 32            (Input, _) => other,
 33            (Witness, _) => *self,
 34        }
 35    }
 36}
 37
 38/// Specifies how variables of type `Self` should be allocated in a
 39/// `ConstraintSystem`.
 40pub trait AllocVar<V: ?Sized, F: Field>: Sized {
 41    /// Allocates a new variable of type `Self` in the `ConstraintSystem` `cs`.
 42    /// The mode of allocation is decided by `mode`.
 43    fn new_variable<T: Borrow<V>>(
 44        cs: impl Into<Namespace<F>>,
 45        f: impl FnOnce() -> Result<T, SynthesisError>,
 46        mode: AllocationMode,
 47    ) -> Result<Self, SynthesisError>;
 48
 49    /// Allocates a new constant of type `Self` in the `ConstraintSystem` `cs`.
 50    ///
 51    /// This should *not* allocate any new variables or constraints in `cs`.
 52    #[tracing::instrument(target = "r1cs", skip(cs, t))]
 53    fn new_constant(
 54        cs: impl Into<Namespace<F>>,
 55        t: impl Borrow<V>,
 56    ) -> Result<Self, SynthesisError> {
 57        Self::new_variable(cs, || Ok(t), AllocationMode::Constant)
 58    }
 59
 60    /// Allocates a new public input of type `Self` in the `ConstraintSystem`
 61    /// `cs`.
 62    #[tracing::instrument(target = "r1cs", skip(cs, f))]
 63    fn new_input<T: Borrow<V>>(
 64        cs: impl Into<Namespace<F>>,
 65        f: impl FnOnce() -> Result<T, SynthesisError>,
 66    ) -> Result<Self, SynthesisError> {
 67        Self::new_variable(cs, f, AllocationMode::Input)
 68    }
 69
 70    /// Allocates a new private witness of type `Self` in the `ConstraintSystem`
 71    /// `cs`.
 72    #[tracing::instrument(target = "r1cs", skip(cs, f))]
 73    fn new_witness<T: Borrow<V>>(
 74        cs: impl Into<Namespace<F>>,
 75        f: impl FnOnce() -> Result<T, SynthesisError>,
 76    ) -> Result<Self, SynthesisError> {
 77        Self::new_variable(cs, f, AllocationMode::Witness)
 78    }
 79}
 80
 81/// This blanket implementation just allocates variables in `Self`
 82/// element by element.
 83impl<I, F: Field, A: AllocVar<I, F>> AllocVar<[I], F> for Vec<A> {
 84    fn new_variable<T: Borrow<[I]>>(
 85        cs: impl Into<Namespace<F>>,
 86        f: impl FnOnce() -> Result<T, SynthesisError>,
 87        mode: AllocationMode,
 88    ) -> Result<Self, SynthesisError> {
 89        let ns = cs.into();
 90        let cs = ns.cs();
 91        f().and_then(|v| {
 92            v.borrow()
 93                .iter()
 94                .map(|e| A::new_variable(cs.clone(), || Ok(e), mode))
 95                .collect()
 96        })
 97    }
 98}
 99
100/// Dummy impl for `()`.
101impl<F: Field> AllocVar<(), F> for () {
102    fn new_variable<T: Borrow<()>>(
103        _cs: impl Into<Namespace<F>>,
104        _f: impl FnOnce() -> Result<T, SynthesisError>,
105        _mode: AllocationMode,
106    ) -> Result<Self, SynthesisError> {
107        Ok(())
108    }
109}
110
111/// This blanket implementation just allocates variables in `Self`
112/// element by element.
113impl<I, F: Field, A: AllocVar<I, F>, const N: usize> AllocVar<[I; N], F> for [A; N] {
114    fn new_variable<T: Borrow<[I; N]>>(
115        cs: impl Into<Namespace<F>>,
116        f: impl FnOnce() -> Result<T, SynthesisError>,
117        mode: AllocationMode,
118    ) -> Result<Self, SynthesisError> {
119        let ns = cs.into();
120        let cs = ns.cs();
121        f().map(|v| {
122            let v = v.borrow();
123            core::array::from_fn(|i| A::new_variable(cs.clone(), || Ok(&v[i]), mode).unwrap())
124        })
125    }
126}
127
128/// This blanket implementation just allocates variables in `Self`
129/// element by element.
130impl<I, F: Field, A: AllocVar<I, F>, const N: usize> AllocVar<[I], F> for [A; N] {
131    fn new_variable<T: Borrow<[I]>>(
132        cs: impl Into<Namespace<F>>,
133        f: impl FnOnce() -> Result<T, SynthesisError>,
134        mode: AllocationMode,
135    ) -> Result<Self, SynthesisError> {
136        let ns = cs.into();
137        let cs = ns.cs();
138        f().map(|v| {
139            let v = v.borrow();
140            core::array::from_fn(|i| A::new_variable(cs.clone(), || Ok(&v[i]), mode).unwrap())
141        })
142    }
143}