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    fn new_constant(
38        cs: impl Into<Namespace<F>>,
39        t: impl Borrow<V>,
40    ) -> Result<Self, SynthesisError> {
41        Self::new_variable(cs, || Ok(t), AllocationMode::Constant)
42    }
43
44    fn new_input<T: Borrow<V>>(
45        cs: impl Into<Namespace<F>>,
46        f: impl FnOnce() -> Result<T, SynthesisError>,
47    ) -> Result<Self, SynthesisError> {
48        Self::new_variable(cs, f, AllocationMode::Input)
49    }
50
51    fn new_witness<T: Borrow<V>>(
52        cs: impl Into<Namespace<F>>,
53        f: impl FnOnce() -> Result<T, SynthesisError>,
54    ) -> Result<Self, SynthesisError> {
55        Self::new_variable(cs, f, AllocationMode::Witness)
56    }
57}
58
59impl<I, F: Field, A: AllocVar<I, F>> AllocVar<[I], F> for Vec<A> {
60    fn new_variable<T: Borrow<[I]>>(
61        cs: impl Into<Namespace<F>>,
62        f: impl FnOnce() -> Result<T, SynthesisError>,
63        mode: AllocationMode,
64    ) -> Result<Self, SynthesisError> {
65        let ns = cs.into();
66        let cs = ns.cs();
67        let mut vec = Vec::new();
68        for value in f()?.borrow().iter() {
69            vec.push(A::new_variable(cs.clone(), || Ok(value), mode)?);
70        }
71        Ok(vec)
72    }
73}