1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaVM library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use crate::r1cs::{ConstraintSystem, Index, LinearCombination, OptionalVec, Variable, errors::SynthesisError};
 17  use alphavm_fields::Field;
 18  
 19  use cfg_if::cfg_if;
 20  use fxhash::{FxBuildHasher, FxHashMap};
 21  use indexmap::{IndexMap, IndexSet, map::Entry};
 22  use itertools::Itertools;
 23  
 24  /// This field is the scalar field (Fr) of BLS12-377.
 25  pub type Fr = alphavm_curves::bls12_377::Fr;
 26  
 27  #[allow(dead_code)]
 28  #[derive(Debug, Clone)]
 29  enum NamedObject {
 30      Constraint(usize),
 31      Var(Variable),
 32      // contains the list of named objects that belong to it
 33      Namespace(Namespace),
 34  }
 35  
 36  #[derive(Debug, Clone, Default)]
 37  struct Namespace {
 38      children: Vec<NamedObject>,
 39      idx: NamespaceIndex,
 40  }
 41  
 42  impl Namespace {
 43      fn push(&mut self, child: NamedObject) {
 44          self.children.push(child);
 45      }
 46  }
 47  
 48  type InternedField = usize;
 49  type InternedPathSegment = usize;
 50  type NamespaceIndex = usize;
 51  
 52  #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 53  pub struct InternedPath {
 54      parent_namespace: NamespaceIndex,
 55      last_segment: InternedPathSegment,
 56  }
 57  
 58  #[derive(PartialEq, Eq, Hash)]
 59  pub struct TestConstraint {
 60      interned_path: InternedPath,
 61      a: Vec<(Variable, InternedField)>,
 62      b: Vec<(Variable, InternedField)>,
 63      c: Vec<(Variable, InternedField)>,
 64  }
 65  
 66  #[derive(Default, Debug)]
 67  pub struct CurrentNamespace {
 68      segments: Vec<InternedPathSegment>,
 69      indices: Vec<NamespaceIndex>,
 70  }
 71  
 72  impl CurrentNamespace {
 73      fn idx(&self) -> usize {
 74          self.indices.last().copied().unwrap_or(0)
 75      }
 76  
 77      fn pop(&mut self) {
 78          assert!(self.segments.pop().is_some());
 79          assert!(self.indices.pop().is_some());
 80      }
 81  }
 82  
 83  /// Constraint system for testing purposes.
 84  pub struct TestConstraintSystem<F: Field> {
 85      // used to intern full paths in test scenarios, for get and set purposes
 86      interned_full_paths: FxHashMap<Vec<InternedPathSegment>, InternedPath>,
 87      // used to intern namespace segments
 88      interned_path_segments: IndexSet<String, FxBuildHasher>,
 89      // used to intern fields belonging to F
 90      interned_fields: IndexSet<F, FxBuildHasher>,
 91      // contains named objects bound to their (interned) paths; the indices are
 92      // used for NamespaceIndex lookups
 93      named_objects: IndexMap<InternedPath, NamedObject, FxBuildHasher>,
 94      // a stack of current path's segments and the index of the current path's
 95      // index in the named_objects map
 96      current_namespace: CurrentNamespace,
 97      // the list of currently applicable constraints
 98      constraints: OptionalVec<TestConstraint>,
 99      // the list of currently applicable input variables
100      public_variables: OptionalVec<InternedField>,
101      // the list of currently applicable auxiliary variables
102      private_variables: OptionalVec<InternedField>,
103  }
104  
105  impl<F: Field> Default for TestConstraintSystem<F> {
106      fn default() -> Self {
107          let mut interned_path_segments = IndexSet::with_hasher(FxBuildHasher::default());
108          let path_segment = "ONE".to_owned();
109          let interned_path_segment = interned_path_segments.insert_full(path_segment).0;
110          let interned_path = InternedPath { parent_namespace: 0, last_segment: interned_path_segment };
111  
112          cfg_if! {
113              if #[cfg(debug_assertions)] {
114                  let mut interned_full_paths = FxHashMap::default();
115                  interned_full_paths.insert(vec![interned_path_segment], interned_path);
116              } else {
117                  let interned_full_paths = FxHashMap::default();
118              }
119          }
120  
121          let mut named_objects = IndexMap::with_hasher(FxBuildHasher::default());
122          named_objects.insert_full(interned_path, NamedObject::Var(TestConstraintSystem::<F>::one()));
123  
124          let mut interned_fields = IndexSet::with_hasher(FxBuildHasher::default());
125          let interned_field = interned_fields.insert_full(F::one()).0;
126  
127          let mut inputs: OptionalVec<InternedField> = Default::default();
128          inputs.insert(interned_field);
129  
130          let constraints = OptionalVec::default();
131  
132          TestConstraintSystem {
133              interned_full_paths,
134              interned_fields,
135              interned_path_segments,
136              named_objects,
137              current_namespace: Default::default(),
138              constraints,
139              public_variables: inputs,
140              private_variables: Default::default(),
141          }
142      }
143  }
144  
145  impl<F: Field> TestConstraintSystem<F> {
146      pub fn new() -> Self {
147          Self::default()
148      }
149  
150      /// Prints the constraint at which `self` and `other` differ.
151      pub fn diff(&self, other: &Self) {
152          for (i, (self_c, other_c)) in self.constraints.iter().zip(other.constraints.iter()).enumerate() {
153              let self_interned_path = self_c.interned_path;
154              let other_interned_path = other_c.interned_path;
155              if self_c.a != other_c.a {
156                  println!("A row {i} is different:");
157                  println!("self: {}", self.unintern_path(self_interned_path));
158                  println!("other: {}", other.unintern_path(other_interned_path));
159                  break;
160              }
161  
162              if self_c.b != other_c.b {
163                  println!("B row {i} is different:");
164                  println!("self: {}", self.unintern_path(self_interned_path));
165                  println!("other: {}", other.unintern_path(other_interned_path));
166                  break;
167              }
168  
169              if self_c.c != other_c.c {
170                  println!("C row {i} is different:");
171                  println!("self: {}", self.unintern_path(self_interned_path));
172                  println!("other: {}", other.unintern_path(other_interned_path));
173                  break;
174              }
175          }
176      }
177  
178      #[inline]
179      fn intern_path(&self, path: &str) -> InternedPath {
180          let mut vec = vec![];
181  
182          for segment in path.split('/') {
183              vec.push(self.interned_path_segments.get_index_of(segment).unwrap());
184          }
185  
186          *self.interned_full_paths.get(&vec).unwrap()
187      }
188  
189      fn unintern_path(&self, interned_path: InternedPath) -> String {
190          let last_segment = self.interned_path_segments.get_index(interned_path.last_segment).unwrap();
191          let mut reversed_uninterned_segments = vec![last_segment];
192  
193          let mut parent_ns = interned_path.parent_namespace;
194          while parent_ns != 0 {
195              let interned_parent_ns = self.named_objects.get_index(parent_ns).unwrap().0;
196              let parent_segment = self.interned_path_segments.get_index(interned_parent_ns.last_segment).unwrap();
197              reversed_uninterned_segments.push(parent_segment);
198              parent_ns = interned_parent_ns.parent_namespace;
199          }
200  
201          let segments = reversed_uninterned_segments.into_iter().map(|s| s.as_str()).rev();
202          Itertools::intersperse(segments, "/").collect()
203      }
204  
205      pub fn print_named_objects(&self) {
206          for TestConstraint { interned_path, .. } in self.constraints.iter() {
207              println!("{}", self.unintern_path(*interned_path));
208          }
209      }
210  
211      fn eval_lc(&self, terms: &[(Variable, InternedField)]) -> F {
212          let mut acc = F::zero();
213  
214          for &(var, interned_coeff) in terms {
215              let interned_tmp = match var.get_unchecked() {
216                  Index::Public(index) => self.public_variables[index],
217                  Index::Private(index) => self.private_variables[index],
218              };
219              let mut tmp = *self.interned_fields.get_index(interned_tmp).unwrap();
220              let coeff = self.interned_fields.get_index(interned_coeff).unwrap();
221  
222              tmp.mul_assign(coeff);
223              acc.add_assign(tmp);
224          }
225  
226          acc
227      }
228  
229      pub fn which_is_unsatisfied(&self) -> Option<String> {
230          for TestConstraint { interned_path, a, b, c } in self.constraints.iter() {
231              let mut a = self.eval_lc(a.as_ref());
232              let b = self.eval_lc(b.as_ref());
233              let c = self.eval_lc(c.as_ref());
234  
235              a.mul_assign(&b);
236  
237              if a != c {
238                  return Some(self.unintern_path(*interned_path));
239              }
240          }
241  
242          None
243      }
244  
245      #[inline]
246      pub fn is_satisfied(&self) -> bool {
247          self.which_is_unsatisfied().is_none()
248      }
249  
250      #[inline]
251      pub fn num_non_zero(&self) -> (usize, usize, usize) {
252          let mut non_zero_a = 0;
253          let mut non_zero_b = 0;
254          let mut non_zero_c = 0;
255          for TestConstraint { a, b, c, .. } in self.constraints.iter() {
256              non_zero_a += a.len();
257              non_zero_b += b.len();
258              non_zero_c += c.len();
259          }
260          (non_zero_a, non_zero_b, non_zero_c)
261      }
262  
263      #[inline]
264      pub fn num_constraints(&self) -> usize {
265          self.constraints.len()
266      }
267  
268      #[inline]
269      pub fn get_constraint_path(&self, i: usize) -> String {
270          self.unintern_path(self.constraints.iter().nth(i).unwrap().interned_path)
271      }
272  
273      pub fn set(&mut self, path: &str, to: F) {
274          let interned_path = self.intern_path(path);
275          let interned_field = self.interned_fields.insert_full(to).0;
276  
277          match self.named_objects.get(&interned_path) {
278              Some(NamedObject::Var(v)) => match v.get_unchecked() {
279                  Index::Public(index) => self.public_variables[index] = interned_field,
280                  Index::Private(index) => self.private_variables[index] = interned_field,
281              },
282              Some(e) => panic!("tried to set path `{path}` to value, but `{e:?}` already exists there."),
283              _ => panic!("no variable exists at path: {path}"),
284          }
285      }
286  
287      pub fn get(&mut self, path: &str) -> F {
288          let interned_path = self.intern_path(path);
289  
290          let interned_field = match self.named_objects.get(&interned_path) {
291              Some(NamedObject::Var(v)) => match v.get_unchecked() {
292                  Index::Public(index) => self.public_variables[index],
293                  Index::Private(index) => self.private_variables[index],
294              },
295              Some(e) => panic!("tried to get value of path `{path}`, but `{e:?}` exists there (not a variable)"),
296              _ => panic!("no variable exists at path: {path}"),
297          };
298  
299          *self.interned_fields.get_index(interned_field).unwrap()
300      }
301  
302      #[inline]
303      fn set_named_obj(&mut self, interned_path: InternedPath, to: NamedObject) -> NamespaceIndex {
304          match self.named_objects.entry(interned_path) {
305              Entry::Vacant(e) => {
306                  let ns_idx = e.index();
307                  e.insert(to);
308                  ns_idx
309              }
310              Entry::Occupied(e) => {
311                  let interned_segments = e.swap_remove_entry().0;
312                  panic!("tried to create object at existing path: {}", self.unintern_path(interned_segments));
313              }
314          }
315      }
316  
317      #[inline]
318      fn compute_path(&mut self, new_segment: &str) -> InternedPath {
319          let (interned_segment, new) = if let Some(index) = self.interned_path_segments.get_index_of(new_segment) {
320              (index, false)
321          } else {
322              self.interned_path_segments.insert_full(new_segment.to_owned())
323          };
324  
325          // only perform the check for segments not seen before
326          assert!(!new || !new_segment.contains('/'), "'/' is not allowed in names");
327  
328          let interned_path =
329              InternedPath { parent_namespace: self.current_namespace.idx(), last_segment: interned_segment };
330  
331          cfg_if! {
332              if #[cfg(debug_assertions)] {
333                  let mut full_path = self.current_namespace.segments.clone();
334                  full_path.push(interned_segment);
335                  self.interned_full_paths.insert(full_path, interned_path);
336              }
337          }
338  
339          interned_path
340      }
341  
342      #[cfg(not(debug_assertions))]
343      fn purge_namespace(&mut self, namespace: Namespace) {
344          for child_obj in namespace.children {
345              match child_obj {
346                  NamedObject::Var(var) => match var.get_unchecked() {
347                      Index::Private(idx) => {
348                          self.private_variables.remove(idx);
349                      }
350                      Index::Public(idx) => {
351                          self.public_variables.remove(idx);
352                      }
353                  },
354                  NamedObject::Constraint(idx) => {
355                      self.constraints.remove(idx);
356                  }
357                  NamedObject::Namespace(children) => {
358                      self.purge_namespace(children);
359                  }
360              }
361              self.named_objects.swap_remove_index(namespace.idx);
362          }
363      }
364  
365      #[inline]
366      fn register_object_in_namespace(&mut self, named_obj: NamedObject) {
367          if let NamedObject::Namespace(ns) = self.named_objects.get_index_mut(self.current_namespace.idx()).unwrap().1 {
368              ns.push(named_obj);
369          }
370      }
371  }
372  
373  impl<F: Field> ConstraintSystem<F> for TestConstraintSystem<F> {
374      type Root = Self;
375  
376      fn alloc<Fn, A, AR>(&mut self, annotation: A, f: Fn) -> Result<Variable, SynthesisError>
377      where
378          Fn: FnOnce() -> Result<F, SynthesisError>,
379          A: FnOnce() -> AR,
380          AR: AsRef<str>,
381      {
382          let interned_path = self.compute_path(annotation().as_ref());
383          let interned_field = self.interned_fields.insert_full(f()?).0;
384          let index = self.private_variables.insert(interned_field);
385          let var = Variable::new_unchecked(Index::Private(index));
386          let named_obj = NamedObject::Var(var);
387          self.register_object_in_namespace(named_obj.clone());
388          self.set_named_obj(interned_path, named_obj);
389  
390          Ok(var)
391      }
392  
393      fn alloc_input<Fn, A, AR>(&mut self, annotation: A, f: Fn) -> Result<Variable, SynthesisError>
394      where
395          Fn: FnOnce() -> Result<F, SynthesisError>,
396          A: FnOnce() -> AR,
397          AR: AsRef<str>,
398      {
399          let interned_path = self.compute_path(annotation().as_ref());
400          let interned_field = self.interned_fields.insert_full(f()?).0;
401          let index = self.public_variables.insert(interned_field);
402          let var = Variable::new_unchecked(Index::Public(index));
403          let named_obj = NamedObject::Var(var);
404          self.register_object_in_namespace(named_obj.clone());
405          self.set_named_obj(interned_path, named_obj);
406  
407          Ok(var)
408      }
409  
410      fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
411      where
412          A: FnOnce() -> AR,
413          AR: AsRef<str>,
414          LA: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
415          LB: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
416          LC: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
417      {
418          let interned_path = self.compute_path(annotation().as_ref());
419          let index = self.constraints.next_idx();
420          let named_obj = NamedObject::Constraint(index);
421          self.register_object_in_namespace(named_obj.clone());
422          self.set_named_obj(interned_path, named_obj);
423  
424          let mut intern_fields = |uninterned: Vec<(Variable, F)>| -> Vec<(Variable, InternedField)> {
425              uninterned
426                  .into_iter()
427                  .map(|(var, field)| {
428                      let interned_field = self.interned_fields.insert_full(field).0;
429                      (var, interned_field)
430                  })
431                  .collect()
432          };
433  
434          let a = intern_fields(a(LinearCombination::zero()).0);
435          let b = intern_fields(b(LinearCombination::zero()).0);
436          let c = intern_fields(c(LinearCombination::zero()).0);
437  
438          self.constraints.insert(TestConstraint { interned_path, a, b, c });
439      }
440  
441      fn push_namespace<NR: AsRef<str>, N: FnOnce() -> NR>(&mut self, name_fn: N) {
442          let name = name_fn();
443          let interned_path = self.compute_path(name.as_ref());
444          let new_segment = interned_path.last_segment;
445          let named_obj = NamedObject::Namespace(Default::default());
446          self.register_object_in_namespace(named_obj.clone());
447          let namespace_idx = self.set_named_obj(interned_path, named_obj);
448          if let NamedObject::Namespace(ref mut ns) = self.named_objects[namespace_idx] {
449              ns.idx = namespace_idx;
450          };
451  
452          self.current_namespace.segments.push(new_segment);
453          self.current_namespace.indices.push(namespace_idx);
454      }
455  
456      #[cfg(not(debug_assertions))]
457      fn pop_namespace(&mut self) {
458          let namespace = if let NamedObject::Namespace(no) =
459              self.named_objects.swap_remove_index(self.current_namespace.idx()).unwrap().1
460          {
461              no
462          } else {
463              unreachable!()
464          };
465  
466          // remove object belonging to the popped namespace
467          self.purge_namespace(namespace);
468  
469          // update the current namespace
470          self.current_namespace.pop();
471      }
472  
473      #[cfg(debug_assertions)]
474      fn pop_namespace(&mut self) {
475          // don't perform a full cleanup in test conditions, so that all the variables,
476          // constraints and namespace indices remain available throughout the tests
477  
478          self.current_namespace.pop();
479      }
480  
481      #[inline]
482      fn get_root(&mut self) -> &mut Self::Root {
483          self
484      }
485  
486      #[inline]
487      fn num_constraints(&self) -> usize {
488          self.num_constraints()
489      }
490  
491      #[inline]
492      fn num_public_variables(&self) -> usize {
493          self.public_variables.len()
494      }
495  
496      #[inline]
497      fn num_private_variables(&self) -> usize {
498          self.private_variables.len()
499      }
500  
501      #[inline]
502      fn is_in_setup_mode(&self) -> bool {
503          false
504      }
505  }