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, Variable, errors::SynthesisError};
 17  use alphavm_fields::Field;
 18  
 19  /// Constraint system for testing purposes.
 20  pub struct TestConstraintChecker<F: Field> {
 21      // the list of currently applicable input variables
 22      public_variables: Vec<F>,
 23      // the list of currently applicable auxiliary variables
 24      private_variables: Vec<F>,
 25      // whether or not unsatisfactory constraint has been found
 26      found_unsatisfactory_constraint: bool,
 27      // number of constraints
 28      num_constraints: usize,
 29      // constraint path segments in the stack
 30      segments: Vec<String>,
 31      // the first unsatisfied constraint
 32      first_unsatisfied_constraint: Option<String>,
 33  }
 34  
 35  impl<F: Field> Default for TestConstraintChecker<F> {
 36      fn default() -> Self {
 37          Self {
 38              public_variables: vec![F::one()],
 39              private_variables: vec![],
 40              found_unsatisfactory_constraint: false,
 41              num_constraints: 0,
 42              segments: vec![],
 43              first_unsatisfied_constraint: None,
 44          }
 45      }
 46  }
 47  
 48  impl<F: Field> TestConstraintChecker<F> {
 49      pub fn new() -> Self {
 50          Self::default()
 51      }
 52  
 53      pub fn which_is_unsatisfied(&self) -> Option<String> {
 54          self.first_unsatisfied_constraint.clone()
 55      }
 56  
 57      #[inline]
 58      pub fn is_satisfied(&self) -> bool {
 59          !self.found_unsatisfactory_constraint
 60      }
 61  
 62      #[inline]
 63      pub fn num_constraints(&self) -> usize {
 64          self.num_constraints
 65      }
 66  
 67      #[inline]
 68      pub fn public_inputs(&self) -> Vec<F> {
 69          self.public_variables[1..].to_vec()
 70      }
 71  }
 72  
 73  impl<F: Field> ConstraintSystem<F> for TestConstraintChecker<F> {
 74      type Root = Self;
 75  
 76      fn alloc<Fn, A, AR>(&mut self, _annotation: A, f: Fn) -> Result<Variable, SynthesisError>
 77      where
 78          Fn: FnOnce() -> Result<F, SynthesisError>,
 79          A: FnOnce() -> AR,
 80          AR: AsRef<str>,
 81      {
 82          let index = self.private_variables.len();
 83          self.private_variables.push(f()?);
 84          let var = Variable::new_unchecked(Index::Private(index));
 85  
 86          Ok(var)
 87      }
 88  
 89      fn alloc_input<Fn, A, AR>(&mut self, _annotation: A, f: Fn) -> Result<Variable, SynthesisError>
 90      where
 91          Fn: FnOnce() -> Result<F, SynthesisError>,
 92          A: FnOnce() -> AR,
 93          AR: AsRef<str>,
 94      {
 95          let index = self.public_variables.len();
 96          self.public_variables.push(f()?);
 97          let var = Variable::new_unchecked(Index::Public(index));
 98  
 99          Ok(var)
100      }
101  
102      fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
103      where
104          A: FnOnce() -> AR,
105          AR: AsRef<str>,
106          LA: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
107          LB: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
108          LC: FnOnce(LinearCombination<F>) -> LinearCombination<F>,
109      {
110          self.num_constraints += 1;
111  
112          let eval_lc = |lc: Vec<(Variable, F)>| -> F {
113              lc.into_iter()
114                  .map(|(var, coeff)| {
115                      let value = match var.get_unchecked() {
116                          Index::Public(index) => self.public_variables[index],
117                          Index::Private(index) => self.private_variables[index],
118                      };
119                      value * coeff
120                  })
121                  .sum::<F>()
122          };
123  
124          let a = eval_lc(a(LinearCombination::zero()).0);
125          let b = eval_lc(b(LinearCombination::zero()).0);
126          let c = eval_lc(c(LinearCombination::zero()).0);
127  
128          if a * b != c && self.first_unsatisfied_constraint.is_none() {
129              self.found_unsatisfactory_constraint = true;
130  
131              let new = annotation().as_ref().to_string();
132              assert!(!new.contains('/'), "'/' is not allowed in names");
133  
134              let mut path = self.segments.clone();
135              path.push(new);
136              self.first_unsatisfied_constraint = Some(path.join("/"));
137          }
138      }
139  
140      fn push_namespace<NR: AsRef<str>, N: FnOnce() -> NR>(&mut self, name_fn: N) {
141          let new = name_fn().as_ref().to_string();
142          assert!(!new.contains('/'), "'/' is not allowed in names");
143  
144          self.segments.push(new)
145      }
146  
147      fn pop_namespace(&mut self) {
148          self.segments.pop();
149      }
150  
151      #[inline]
152      fn get_root(&mut self) -> &mut Self::Root {
153          self
154      }
155  
156      #[inline]
157      fn num_constraints(&self) -> usize {
158          self.num_constraints()
159      }
160  
161      #[inline]
162      fn num_public_variables(&self) -> usize {
163          self.public_variables.len()
164      }
165  
166      #[inline]
167      fn num_private_variables(&self) -> usize {
168          self.private_variables.len()
169      }
170  
171      #[inline]
172      fn is_in_setup_mode(&self) -> bool {
173          false
174      }
175  }