From 723f197d0e954a72de2eb355a62b9f6a484356f9 Mon Sep 17 00:00:00 2001
From: Robin Salen <salenrobin@gmail.com>
Date: Thu, 20 Apr 2023 08:36:54 +0200
Subject: [PATCH] Cleanup

---
 plonky2/src/gates/gate.rs    |  3 ++-
 plonky2/src/iop/generator.rs | 29 ++++++-----------------------
 2 files changed, 8 insertions(+), 24 deletions(-)

diff --git a/plonky2/src/gates/gate.rs b/plonky2/src/gates/gate.rs
index 9821106c..43c1179a 100644
--- a/plonky2/src/gates/gate.rs
+++ b/plonky2/src/gates/gate.rs
@@ -198,6 +198,7 @@ pub trait Gate<F: RichField + Extendable<D>, const D: usize>: 'static + Send + S
     }
 }
 
+/// A wrapper trait over a `Gate`, to allow for gate serialization.
 pub trait AnyGate<F: RichField + Extendable<D>, const D: usize>: Gate<F, D> {
     fn as_any(&self) -> &dyn Any;
 }
@@ -208,7 +209,7 @@ impl<T: Gate<F, D>, F: RichField + Extendable<D>, const D: usize> AnyGate<F, D>
     }
 }
 
-/// A wrapper around an `Rc<Gate>` which implements `PartialEq`, `Eq` and `Hash` based on gate IDs.
+/// A wrapper around an `Arc<AnyGate>` which implements `PartialEq`, `Eq` and `Hash` based on gate IDs.
 #[derive(Clone)]
 pub struct GateRef<F: RichField + Extendable<D>, const D: usize>(pub(crate) Arc<dyn AnyGate<F, D>>);
 
diff --git a/plonky2/src/iop/generator.rs b/plonky2/src/iop/generator.rs
index 165b5ada..16a23e26 100644
--- a/plonky2/src/iop/generator.rs
+++ b/plonky2/src/iop/generator.rs
@@ -1,6 +1,5 @@
 use alloc::vec;
 use alloc::vec::Vec;
-use core::any::Any;
 use core::fmt::Debug;
 use core::marker::PhantomData;
 
@@ -118,33 +117,19 @@ pub trait WitnessGenerator<F: Field>: 'static + Send + Sync + Debug {
         Self: Sized;
 }
 
-pub trait AnyWitnessGenerator<F: Field>: WitnessGenerator<F> {
-    fn as_any(&self) -> &dyn Any;
-}
-
-impl<T: WitnessGenerator<F>, F: Field> AnyWitnessGenerator<F> for T {
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-}
-
-/// A wrapper around an `Box<AnyWitnessGenerator>`.
-pub struct WitnessGeneratorRef<F: Field>(pub Box<dyn AnyWitnessGenerator<F>>);
+/// A wrapper around an `Box<WitnessGenerator>` which implements `PartialEq`
+/// and `Eq` based on generator IDs.
+pub struct WitnessGeneratorRef<F: Field>(pub Box<dyn WitnessGenerator<F>>);
 
 impl<F: Field> WitnessGeneratorRef<F> {
-    pub fn new<G: AnyWitnessGenerator<F>>(generator: G) -> WitnessGeneratorRef<F> {
+    pub fn new<G: WitnessGenerator<F>>(generator: G) -> WitnessGeneratorRef<F> {
         WitnessGeneratorRef(Box::new(generator))
     }
 }
 
 impl<F: Field> PartialEq for WitnessGeneratorRef<F> {
     fn eq(&self, other: &Self) -> bool {
-        let mut buf1 = Vec::new();
-        let mut buf2 = Vec::new();
-        self.0.serialize(&mut buf1).unwrap();
-        other.0.serialize(&mut buf2).unwrap();
-
-        buf1 == buf2
+        self.0.id() == other.0.id()
     }
 }
 
@@ -152,9 +137,7 @@ impl<F: Field> Eq for WitnessGeneratorRef<F> {}
 
 impl<F: Field> Debug for WitnessGeneratorRef<F> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        let mut buf = Vec::new();
-        self.0.serialize(&mut buf).unwrap();
-        write!(f, "{:?}", buf)
+        write!(f, "{}", self.0.id())
     }
 }