1use algebra::Field;
  2use r1cs_core::{ConstraintSystem, Index, LinearCombination, SynthesisError, Variable};
  3
  4use radix_trie::Trie;
  5
  6#[derive(Debug)]
  7enum NamedObject {
  8    Constraint(usize),
  9    Var(Variable),
 10    Namespace,
 11}
 12
 13/// Constraint system for testing purposes.
 14pub struct TestConstraintSystem<ConstraintF: Field> {
 15    named_objects:     Trie<String, NamedObject>,
 16    current_namespace: Vec<String>,
 17    pub constraints: Vec<(
 18        LinearCombination<ConstraintF>,
 19        LinearCombination<ConstraintF>,
 20        LinearCombination<ConstraintF>,
 21        String,
 22    )>,
 23    inputs:            Vec<(ConstraintF, String)>,
 24    aux:               Vec<(ConstraintF, String)>,
 25}
 26
 27impl<ConstraintF: Field> TestConstraintSystem<ConstraintF> {
 28    fn eval_lc(
 29        terms: &[(Variable, ConstraintF)],
 30        inputs: &[(ConstraintF, String)],
 31        aux: &[(ConstraintF, String)],
 32    ) -> ConstraintF {
 33        let mut acc = ConstraintF::zero();
 34
 35        for &(var, ref coeff) in terms {
 36            let mut tmp = match var.get_unchecked() {
 37                Index::Input(index) => inputs[index].0,
 38                Index::Aux(index) => aux[index].0,
 39            };
 40
 41            tmp.mul_assign(&coeff);
 42            acc.add_assign(&tmp);
 43        }
 44
 45        acc
 46    }
 47}
 48
 49impl<ConstraintF: Field> TestConstraintSystem<ConstraintF> {
 50    pub fn new() -> TestConstraintSystem<ConstraintF> {
 51        let mut map = Trie::new();
 52        map.insert(
 53            "ONE".into(),
 54            NamedObject::Var(TestConstraintSystem::<ConstraintF>::one()),
 55        );
 56
 57        TestConstraintSystem {
 58            named_objects:     map,
 59            current_namespace: vec![],
 60            constraints:       vec![],
 61            inputs:            vec![(ConstraintF::one(), "ONE".into())],
 62            aux:               vec![],
 63        }
 64    }
 65
 66    pub fn print_named_objects(&self) {
 67        for &(_, _, _, ref name) in &self.constraints {
 68            println!("{}", name);
 69        }
 70    }
 71
 72    pub fn which_is_unsatisfied(&self) -> Option<&str> {
 73        for &(ref a, ref b, ref c, ref path) in &self.constraints {
 74            let mut a = Self::eval_lc(a.as_ref(), &self.inputs, &self.aux);
 75            let b = Self::eval_lc(b.as_ref(), &self.inputs, &self.aux);
 76            let c = Self::eval_lc(c.as_ref(), &self.inputs, &self.aux);
 77
 78            a.mul_assign(&b);
 79
 80            if a != c {
 81                return Some(&*path);
 82            }
 83        }
 84
 85        None
 86    }
 87
 88    pub fn is_satisfied(&self) -> bool {
 89        self.which_is_unsatisfied().is_none()
 90    }
 91
 92    pub fn num_constraints(&self) -> usize {
 93        self.constraints.len()
 94    }
 95
 96    pub fn set(&mut self, path: &str, to: ConstraintF) {
 97        match self.named_objects.get(path) {
 98            Some(&NamedObject::Var(ref v)) => match v.get_unchecked() {
 99                Index::Input(index) => self.inputs[index].0 = to,
100                Index::Aux(index) => self.aux[index].0 = to,
101            },
102            Some(e) => panic!(
103                "tried to set path `{}` to value, but `{:?}` already exists there.",
104                path, e
105            ),
106            _ => panic!("no variable exists at path: {}", path),
107        }
108    }
109
110    pub fn get(&mut self, path: &str) -> ConstraintF {
111        match self.named_objects.get(path) {
112            Some(&NamedObject::Var(ref v)) => match v.get_unchecked() {
113                Index::Input(index) => self.inputs[index].0,
114                Index::Aux(index) => self.aux[index].0,
115            },
116            Some(e) => panic!(
117                "tried to get value of path `{}`, but `{:?}` exists there (not a variable)",
118                path, e
119            ),
120            _ => panic!("no variable exists at path: {}", path),
121        }
122    }
123
124    fn set_named_obj(&mut self, path: String, to: NamedObject) {
125        if self.named_objects.get(&path).is_some() {
126            panic!("tried to create object at existing path: {}", path);
127        }
128
129        self.named_objects.insert(path, to);
130    }
131}
132
133fn compute_path(ns: &[String], this: String) -> String {
134    if this.chars().any(|a| a == '/') {
135        panic!("'/' is not allowed in names");
136    }
137
138    let mut name = String::new();
139
140    let mut needs_separation = false;
141    for ns in ns.iter().chain(Some(&this).into_iter()) {
142        if needs_separation {
143            name += "/";
144        }
145
146        name += ns;
147        needs_separation = true;
148    }
149
150    name
151}
152
153impl<ConstraintF: Field> ConstraintSystem<ConstraintF> for TestConstraintSystem<ConstraintF> {
154    type Root = Self;
155
156    fn alloc<F, A, AR>(&mut self, annotation: A, f: F) -> Result<Variable, SynthesisError>
157    where
158        F: FnOnce() -> Result<ConstraintF, SynthesisError>,
159        A: FnOnce() -> AR,
160        AR: Into<String>,
161    {
162        let index = self.aux.len();
163        let path = compute_path(&self.current_namespace, annotation().into());
164        self.aux.push((f()?, path.clone()));
165        let var = Variable::new_unchecked(Index::Aux(index));
166        self.set_named_obj(path, NamedObject::Var(var));
167
168        Ok(var)
169    }
170
171    fn alloc_input<F, A, AR>(&mut self, annotation: A, f: F) -> Result<Variable, SynthesisError>
172    where
173        F: FnOnce() -> Result<ConstraintF, SynthesisError>,
174        A: FnOnce() -> AR,
175        AR: Into<String>,
176    {
177        let index = self.inputs.len();
178        let path = compute_path(&self.current_namespace, annotation().into());
179        self.inputs.push((f()?, path.clone()));
180        let var = Variable::new_unchecked(Index::Input(index));
181        self.set_named_obj(path, NamedObject::Var(var));
182
183        Ok(var)
184    }
185
186    fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
187    where
188        A: FnOnce() -> AR,
189        AR: Into<String>,
190        LA: FnOnce(LinearCombination<ConstraintF>) -> LinearCombination<ConstraintF>,
191        LB: FnOnce(LinearCombination<ConstraintF>) -> LinearCombination<ConstraintF>,
192        LC: FnOnce(LinearCombination<ConstraintF>) -> LinearCombination<ConstraintF>,
193    {
194        let path = compute_path(&self.current_namespace, annotation().into());
195        let index = self.constraints.len();
196        self.set_named_obj(path.clone(), NamedObject::Constraint(index));
197
198        let mut a = a(LinearCombination::zero());
199        let mut b = b(LinearCombination::zero());
200        let mut c = c(LinearCombination::zero());
201        a.0.shrink_to_fit();
202        b.0.shrink_to_fit();
203        c.0.shrink_to_fit();
204
205        self.constraints.push((a, b, c, path));
206    }
207
208    fn push_namespace<NR, N>(&mut self, name_fn: N)
209    where
210        NR: Into<String>,
211        N: FnOnce() -> NR,
212    {
213        let name = name_fn().into();
214        let path = compute_path(&self.current_namespace, name.clone());
215        self.set_named_obj(path, NamedObject::Namespace);
216        self.current_namespace.push(name);
217    }
218
219    fn pop_namespace(&mut self) {
220        assert!(self.current_namespace.pop().is_some());
221    }
222
223    fn get_root(&mut self) -> &mut Self::Root {
224        self
225    }
226
227    fn num_constraints(&self) -> usize {
228        self.constraints.len()
229    }
230}