From 801563369b1a2f51045a1890d11bfa2770275f8b Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 24 Feb 2021 13:07:22 -0800
Subject: [PATCH] Mostly finish GMiMC gate

---
 src/constraint_polynomial.rs | 52 ++++++++++++++++++++++++++++++++++--
 src/gates/gmimc.rs           | 24 +++++++++++++++--
 2 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/src/constraint_polynomial.rs b/src/constraint_polynomial.rs
index c569bc08..cf32f619 100644
--- a/src/constraint_polynomial.rs
+++ b/src/constraint_polynomial.rs
@@ -1,7 +1,7 @@
 use std::collections::{HashMap, HashSet};
 use std::hash::{Hash, Hasher};
 use std::iter::{Product, Sum};
-use std::ops::{Add, Mul, Neg, Sub};
+use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
 use std::ptr;
 use std::rc::Rc;
 
@@ -101,6 +101,10 @@ impl<F: Field> ConstraintPolynomial<F> {
         self * self
     }
 
+    pub fn cube(&self) -> Self {
+        self * self * self
+    }
+
     pub(crate) fn degree(&self) -> usize {
         (self.0).0.degree()
     }
@@ -181,7 +185,7 @@ impl<F: Field> Neg for &ConstraintPolynomial<F> {
 /// Takes the following arguments:
 /// - `$trait`: the name of the binary operation trait to implement
 /// - `$method`: the name of the method in the trait. It is assumed that `ConstraintPolynomial`
-///   contains a method with the same name, implementing the `Self . Self` variant.
+///   contains a method with the same name, implementing the `&Self . &Self` variant.
 macro_rules! binop_variants {
     ($trait:ident, $method:ident) => {
         impl<F: Field> $trait<Self> for ConstraintPolynomial<F> {
@@ -250,10 +254,54 @@ macro_rules! binop_variants {
     };
 }
 
+/// Generates the following variants of a binary operation assignment:
+/// - `.= Self`
+/// - `.= &Self`
+/// - `.= F`
+/// - `.= usize`
+/// where `Self` is `ConstraintPolynomial<F>`.
+///
+/// Takes the following arguments:
+/// - `$trait`: the name of the binary operation trait to implement
+/// - `$assign_method`: the name of the method in the trait
+/// - `$binop_method`: the name of a method in `ConstraintPolynomial`
+///   which implements the `&Self . &Self` operation.
+macro_rules! binop_assign_variants {
+    ($trait:ident, $assign_method:ident, $binop_method:ident) => {
+        impl<F: Field> $trait<Self> for ConstraintPolynomial<F> {
+            fn $assign_method(&mut self, rhs: Self) {
+                *self = ConstraintPolynomial::$binop_method(self, &rhs);
+            }
+        }
+
+        impl<F: Field> $trait<&Self> for ConstraintPolynomial<F> {
+            fn $assign_method(&mut self, rhs: &Self) {
+                *self = ConstraintPolynomial::$binop_method(self, rhs);
+            }
+        }
+
+        impl<F: Field> $trait<F> for ConstraintPolynomial<F> {
+            fn $assign_method(&mut self, rhs: F) {
+                *self = ConstraintPolynomial::$binop_method(self, &ConstraintPolynomial::constant(rhs));
+            }
+        }
+
+        impl<F: Field> $trait<usize> for ConstraintPolynomial<F> {
+            fn $assign_method(&mut self, rhs: usize) {
+                *self = ConstraintPolynomial::$binop_method(self, &ConstraintPolynomial::constant_usize(rhs));
+            }
+        }
+    }
+}
+
 binop_variants!(Add, add);
 binop_variants!(Sub, sub);
 binop_variants!(Mul, mul);
 
+binop_assign_variants!(AddAssign, add_assign, add);
+binop_assign_variants!(SubAssign, sub_assign, sub);
+binop_assign_variants!(MulAssign, mul_assign, mul);
+
 impl<F: Field> Sum for ConstraintPolynomial<F> {
     fn sum<I: Iterator<Item=Self>>(iter: I) -> Self {
         iter.fold(
diff --git a/src/gates/gmimc.rs b/src/gates/gmimc.rs
index 79acc8b4..f76b6410 100644
--- a/src/gates/gmimc.rs
+++ b/src/gates/gmimc.rs
@@ -1,4 +1,5 @@
 use std::convert::TryInto;
+use std::sync::Arc;
 
 use crate::circuit_data::CircuitConfig;
 use crate::constraint_polynomial::ConstraintPolynomial;
@@ -10,7 +11,6 @@ use crate::gmimc::gmimc_permute;
 use crate::target::Target2;
 use crate::wire::Wire;
 use crate::witness::PartialWitness;
-use std::sync::Arc;
 
 /// Evaluates a full GMiMC permutation, and writes the output to the next gate's first `width`
 /// wires (which could be the input of another `GMiMCGate`).
@@ -32,7 +32,27 @@ impl<F: Field, const W: usize, const R: usize> Gate<F> for GMiMCGate<F, W, R> {
     }
 
     fn constraints(&self, config: CircuitConfig) -> Vec<ConstraintPolynomial<F>> {
-        unimplemented!()
+        let mut state = (0..W)
+            .map(|i| ConstraintPolynomial::local_wire_value(i))
+            .collect::<Vec<_>>();
+
+        // Value that is implicitly added to each element.
+        // See https://affine.group/2020/02/starkware-challenge
+        let mut addition_buffer = ConstraintPolynomial::zero();
+
+        for r in 0..R {
+            let active = r % W;
+            let round_constant = ConstraintPolynomial::constant(self.round_constants[r]);
+            let f = (&state[active] + &addition_buffer + round_constant).cube();
+            addition_buffer += &f;
+            state[active] -= f;
+        }
+
+        for i in 0..W {
+            state[i] += &addition_buffer;
+        }
+
+        state
     }
 
     fn generators(&self, config: CircuitConfig, gate_index: usize, local_constants: Vec<F>, next_constants: Vec<F>) -> Vec<Box<dyn WitnessGenerator2<F>>> {