From 402dac423da4259f97788827848ee0d9443454d2 Mon Sep 17 00:00:00 2001
From: mmagician <marcin.gorny.94@protonmail.com>
Date: Mon, 7 Aug 2023 19:25:08 -0600
Subject: [PATCH] Add `sub_mul` to GateInstructions (#102)

* Add `sub_mul` to GateInstructions

* Add `sub_mul` prop test
---
 halo2-base/src/gates/flex_gate.rs       | 22 ++++++++++++++++++++++
 halo2-base/src/gates/tests/flex_gate.rs |  9 +++++++++
 halo2-base/src/gates/tests/pos_prop.rs  |  7 +++++++
 halo2-base/src/gates/tests/utils.rs     |  4 ++++
 4 files changed, 42 insertions(+)

diff --git a/halo2-base/src/gates/flex_gate.rs b/halo2-base/src/gates/flex_gate.rs
index 02c7225c..a0447ae7 100644
--- a/halo2-base/src/gates/flex_gate.rs
+++ b/halo2-base/src/gates/flex_gate.rs
@@ -202,6 +202,28 @@ pub trait GateInstructions<F: ScalarField> {
         ctx.get(-4)
     }
 
+    /// Constrains and returns  `a - b * c = out`.
+    ///
+    /// Defines a vertical gate of form | a - b * c | b | c | a |, where (a - b * c) = out.
+    /// * `ctx`: [Context] to add the constraints to
+    /// * `a`: [QuantumCell] value to subtract 'b * c' from
+    /// * `b`: [QuantumCell] value
+    /// * `c`: [QuantumCell] value
+    fn sub_mul(
+        &self,
+        ctx: &mut Context<F>,
+        a: impl Into<QuantumCell<F>>,
+        b: impl Into<QuantumCell<F>>,
+        c: impl Into<QuantumCell<F>>,
+    ) -> AssignedValue<F> {
+        let a = a.into();
+        let b = b.into();
+        let c = c.into();
+        let out_val = *a.value() - *b.value() * c.value();
+        ctx.assign_region_last([Witness(out_val), b, c, a], [0]);
+        ctx.get(-4)
+    }
+
     /// Constrains and returns `a * (-1) = out`.
     ///
     /// Defines a vertical gate of form | a | -a | 1 | 0 |, where (-a) = out.
diff --git a/halo2-base/src/gates/tests/flex_gate.rs b/halo2-base/src/gates/tests/flex_gate.rs
index 8c1689e3..80ef2430 100644
--- a/halo2-base/src/gates/tests/flex_gate.rs
+++ b/halo2-base/src/gates/tests/flex_gate.rs
@@ -31,6 +31,15 @@ pub fn test_sub<F: ScalarField>(inputs: &[QuantumCell<F>]) -> F {
     *a.value()
 }
 
+#[test_case(&[1, 1, 1].map(Fr::from).map(Witness) => Fr::from(0) ; "sub_mul(): 1 - 1 * 1 == 0")]
+pub fn test_sub_mul<F: ScalarField>(inputs: &[QuantumCell<F>]) -> F {
+    let mut builder = GateThreadBuilder::mock();
+    let ctx = builder.main(0);
+    let chip = GateChip::default();
+    let a = chip.sub_mul(ctx, inputs[0], inputs[1], inputs[2]);
+    *a.value()
+}
+
 #[test_case(Witness(Fr::from(1)) => -Fr::from(1) ; "neg(): 1 -> -1")]
 pub fn test_neg<F: ScalarField>(a: QuantumCell<F>) -> F {
     let mut builder = GateThreadBuilder::mock();
diff --git a/halo2-base/src/gates/tests/pos_prop.rs b/halo2-base/src/gates/tests/pos_prop.rs
index 923a76cc..fd79e33a 100644
--- a/halo2-base/src/gates/tests/pos_prop.rs
+++ b/halo2-base/src/gates/tests/pos_prop.rs
@@ -90,6 +90,13 @@ proptest! {
         prop_assert_eq!(result, ground_truth);
     }
 
+    #[test]
+    fn prop_test_sub_mul(input in vec(rand_witness(), 3)) {
+        let ground_truth = sub_mul_ground_truth(input.as_slice());
+        let result = flex_gate::test_sub_mul(input.as_slice());
+        prop_assert_eq!(result, ground_truth);
+    }
+
     #[test]
     fn prop_test_neg(input in rand_witness()) {
         let ground_truth = neg_ground_truth(input);
diff --git a/halo2-base/src/gates/tests/utils.rs b/halo2-base/src/gates/tests/utils.rs
index 234cf636..34e2a435 100644
--- a/halo2-base/src/gates/tests/utils.rs
+++ b/halo2-base/src/gates/tests/utils.rs
@@ -19,6 +19,10 @@ pub fn sub_ground_truth<F: ScalarField>(inputs: &[QuantumCell<F>]) -> F {
     *inputs[0].value() - *inputs[1].value()
 }
 
+pub fn sub_mul_ground_truth<F: ScalarField>(inputs: &[QuantumCell<F>]) -> F {
+    *inputs[0].value() - *inputs[1].value() * *inputs[2].value()
+}
+
 pub fn neg_ground_truth<F: ScalarField>(input: QuantumCell<F>) -> F {
     -(*input.value())
 }