1use crate::Vec;
 2use algebra::Field;
 3use core::borrow::Borrow;
 4use r1cs_core::{Namespace, SynthesisError};
 5
 6#[derive(Eq, PartialEq, Ord, PartialOrd, Debug, Copy, Clone)]
 7pub enum AllocationMode {
 8    Constant = 0,
 9    Input = 1,
10    Witness = 2,
11}
12
13impl AllocationMode {
14    // Outputs the maximum according to the relation `Constant < Input < Witness`.
15    pub fn max(&self, other: Self) -> Self {
16        use AllocationMode::*;
17        match (self, other) {
18            (Constant, _) => other,
19            (Input, Constant) => *self,
20            (Input, _) => other,
21            (Witness, _) => *self,
22        }
23    }
24}
25
26pub trait AllocVar<V, F: Field>
27where
28    Self: Sized,
29    V: ?Sized,
30{
31    fn new_variable<T: Borrow<V>>(
32        cs: impl Into<Namespace<F>>,
33        f: impl FnOnce() -> Result<T, SynthesisError>,
34        mode: AllocationMode,
35    ) -> Result<Self, SynthesisError>;
36
37    #[tracing::instrument(target = "r1cs", skip(cs, t))]
38    fn new_constant(
39        cs: impl Into<Namespace<F>>,
40        t: impl Borrow<V>,
41    ) -> Result<Self, SynthesisError> {
42        Self::new_variable(cs, || Ok(t), AllocationMode::Constant)
43    }
44
45    #[tracing::instrument(target = "r1cs", skip(cs, f))]
46    fn new_input<T: Borrow<V>>(
47        cs: impl Into<Namespace<F>>,
48        f: impl FnOnce() -> Result<T, SynthesisError>,
49    ) -> Result<Self, SynthesisError> {
50        Self::new_variable(cs, f, AllocationMode::Input)
51    }
52
53    #[tracing::instrument(target = "r1cs", skip(cs, f))]
54    fn new_witness<T: Borrow<V>>(
55        cs: impl Into<Namespace<F>>,
56        f: impl FnOnce() -> Result<T, SynthesisError>,
57    ) -> Result<Self, SynthesisError> {
58        Self::new_variable(cs, f, AllocationMode::Witness)
59    }
60}
61
62impl<I, F: Field, A: AllocVar<I, F>> AllocVar<[I], F> for Vec<A> {
63    fn new_variable<T: Borrow<[I]>>(
64        cs: impl Into<Namespace<F>>,
65        f: impl FnOnce() -> Result<T, SynthesisError>,
66        mode: AllocationMode,
67    ) -> Result<Self, SynthesisError> {
68        let ns = cs.into();
69        let cs = ns.cs();
70        let mut vec = Vec::new();
71        for value in f()?.borrow().iter() {
72            vec.push(A::new_variable(cs.clone(), || Ok(value), mode)?);
73        }
74        Ok(vec)
75    }
76}