1use crate::Vec;
 2use algebra::Field;
 3use core::borrow::Borrow;
 4use r1cs_core::{Namespace, SynthesisError};
 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 < Witness`.
26    pub fn max(&self, other: Self) -> Self {
27        use AllocationMode::*;
28        match (self, other) {
29            (Constant, _) => other,
30            (Input, Constant) => *self,
31            (Input, _) => other,
32            (Witness, _) => *self,
33        }
34    }
35}
36
37/// Specifies how variables of type `Self` should be allocated in a `ConstraintSystem`.
38pub trait AllocVar<V, F: Field>
39where
40    Self: Sized,
41    V: ?Sized,
42{
43    /// Allocates a new variable of type `Self` in the `ConstraintSystem` `cs`.
44    /// The mode of allocation is decided by `mode`.
45    fn new_variable<T: Borrow<V>>(
46        cs: impl Into<Namespace<F>>,
47        f: impl FnOnce() -> Result<T, SynthesisError>,
48        mode: AllocationMode,
49    ) -> Result<Self, SynthesisError>;
50
51    /// Allocates a new constant of type `Self` in the `ConstraintSystem` `cs`.
52    ///
53    /// This should *not* allocate any new variables or constraints in `cs`.
54    #[tracing::instrument(target = "r1cs", skip(cs, t))]
55    fn new_constant(
56        cs: impl Into<Namespace<F>>,
57        t: impl Borrow<V>,
58    ) -> Result<Self, SynthesisError> {
59        Self::new_variable(cs, || Ok(t), AllocationMode::Constant)
60    }
61
62    /// Allocates a new public input of type `Self` in the `ConstraintSystem` `cs`.
63    #[tracing::instrument(target = "r1cs", skip(cs, f))]
64    fn new_input<T: Borrow<V>>(
65        cs: impl Into<Namespace<F>>,
66        f: impl FnOnce() -> Result<T, SynthesisError>,
67    ) -> Result<Self, SynthesisError> {
68        Self::new_variable(cs, f, AllocationMode::Input)
69    }
70
71    /// Allocates a new private witness of type `Self` in the `ConstraintSystem` `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        let mut vec = Vec::new();
92        for value in f()?.borrow().iter() {
93            vec.push(A::new_variable(cs.clone(), || Ok(value), mode)?);
94        }
95        Ok(vec)
96    }
97}