From 029eaad0f16a59e603e8d6babf60ad9034569636 Mon Sep 17 00:00:00 2001
From: Varun Thakore <vrnthakore@gmail.com>
Date: Sun, 22 Oct 2023 21:07:43 +0530
Subject: [PATCH] Add an AllocatedNum addition constraint (#39)

* Add AllocatedNum addition and test

* Update addition test
---
 crates/bellpepper-core/src/gadgets/num.rs | 53 +++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/crates/bellpepper-core/src/gadgets/num.rs b/crates/bellpepper-core/src/gadgets/num.rs
index a4b3686..dfe0e6b 100644
--- a/crates/bellpepper-core/src/gadgets/num.rs
+++ b/crates/bellpepper-core/src/gadgets/num.rs
@@ -276,6 +276,38 @@ impl<Scalar: PrimeField> AllocatedNum<Scalar> {
         Ok(bits.into_iter().map(Boolean::from).collect())
     }
 
+    pub fn add<CS>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError>
+    where
+        CS: ConstraintSystem<Scalar>,
+    {
+        let mut value = None;
+
+        let var = cs.alloc(
+            || "sum num",
+            || {
+                let mut tmp = self.value.ok_or(SynthesisError::AssignmentMissing)?;
+                tmp.add_assign(other.value.ok_or(SynthesisError::AssignmentMissing)?);
+
+                value = Some(tmp);
+
+                Ok(tmp)
+            },
+        )?;
+
+        // Constrain: (a + b) * 1 = a + b
+        cs.enforce(
+            || "addition constraint",
+            |lc| lc + self.variable + other.variable,
+            |lc| lc + CS::one(),
+            |lc| lc + var,
+        );
+
+        Ok(AllocatedNum {
+            value,
+            variable: var,
+        })
+    }
+
     pub fn mul<CS>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError>
     where
         CS: ConstraintSystem<Scalar>,
@@ -539,6 +571,27 @@ mod test {
         assert!(cs.get("num") == Fr::ONE);
     }
 
+    #[test]
+    fn test_num_addition() {
+        let mut cs = TestConstraintSystem::<Fr>::new();
+
+        let mut char = Fr::char();
+        char[0] -= 1u8;
+        let mod_minus_one = Fr::from_repr(char);
+        assert!(bool::from(mod_minus_one.is_some()));
+        let mod_minus_one = mod_minus_one.unwrap();
+
+        let a = AllocatedNum::alloc(cs.namespace(|| "a"), || Ok(mod_minus_one)).unwrap();
+        let b = AllocatedNum::alloc(cs.namespace(|| "b"), || Ok(Fr::ONE)).unwrap();
+        let c = a.add(&mut cs, &b).unwrap();
+
+        assert!(cs.is_satisfied());
+        assert!(cs.get("sum num") == Fr::ZERO);
+        assert!(c.value.unwrap() == Fr::ZERO);
+        cs.set("sum num", Fr::ONE);
+        assert!(!cs.is_satisfied());
+    }
+
     #[test]
     fn test_num_squaring() {
         let mut cs = TestConstraintSystem::<Fr>::new();