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`.
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`.
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    #[tracing::instrument(target = "r1cs", skip(cs, f))]
62    fn new_input<T: Borrow<V>>(
63        cs: impl Into<Namespace<F>>,
64        f: impl FnOnce() -> Result<T, SynthesisError>,
65    ) -> Result<Self, SynthesisError> {
66        Self::new_variable(cs, f, AllocationMode::Input)
67    }
68
69    /// Allocates a new private witness of type `Self` in the `ConstraintSystem`.
70    #[tracing::instrument(target = "r1cs", skip(cs, f))]
71    fn new_witness<T: Borrow<V>>(
72        cs: impl Into<Namespace<F>>,
73        f: impl FnOnce() -> Result<T, SynthesisError>,
74    ) -> Result<Self, SynthesisError> {
75        Self::new_variable(cs, f, AllocationMode::Witness)
76    }
77}
78
79impl<I, F: Field, A: AllocVar<I, F>> AllocVar<[I], F> for Vec<A> {
80    fn new_variable<T: Borrow<[I]>>(
81        cs: impl Into<Namespace<F>>,
82        f: impl FnOnce() -> Result<T, SynthesisError>,
83        mode: AllocationMode,
84    ) -> Result<Self, SynthesisError> {
85        let ns = cs.into();
86        let cs = ns.cs();
87        let mut vec = Vec::new();
88        for value in f()?.borrow().iter() {
89            vec.push(A::new_variable(cs.clone(), || Ok(value), mode)?);
90        }
91        Ok(vec)
92    }
93}