From 4af8017fa6921c14080dab7790f519cd9394a7d5 Mon Sep 17 00:00:00 2001
From: Piotr Roslaniec <p.roslaniec@gmail.com>
Date: Mon, 29 Jan 2024 14:56:43 +0100
Subject: [PATCH] refactor: unify share creating methods

---
 ferveo-wasm/tests/node.rs     |  2 +-
 ferveo/src/api.rs             | 12 ++++++++----
 ferveo/src/bindings_python.rs |  4 ++--
 ferveo/src/bindings_wasm.rs   |  2 +-
 ferveo/src/dkg.rs             | 16 ++++++----------
 5 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/ferveo-wasm/tests/node.rs b/ferveo-wasm/tests/node.rs
index 68e5f641..bae5750b 100644
--- a/ferveo-wasm/tests/node.rs
+++ b/ferveo-wasm/tests/node.rs
@@ -32,7 +32,7 @@ fn setup_dkg(shares_num: u32, security_threshold: u32) -> TestSetup {
     // Each validator holds their own DKG instance and generates a transcript every
     // validator, including themselves
     let messages = validators.iter().map(|sender| {
-        let dkg = Dkg::new(
+        let mut dkg = Dkg::new(
             TAU,
             shares_num,
             security_threshold,
diff --git a/ferveo/src/api.rs b/ferveo/src/api.rs
index 26da8c57..aab043d6 100644
--- a/ferveo/src/api.rs
+++ b/ferveo/src/api.rs
@@ -30,7 +30,7 @@ use crate::bindings_python;
 use crate::bindings_wasm;
 pub use crate::EthereumAddress;
 use crate::{
-    do_verify_aggregation, Error, PVSSMap, PubliclyVerifiableParams,
+    do_verify_aggregation, Error, Message, PVSSMap, PubliclyVerifiableParams,
     PubliclyVerifiableSS, Result,
 };
 
@@ -222,10 +222,14 @@ impl Dkg {
     }
 
     pub fn generate_transcript<R: RngCore>(
-        &self,
+        &mut self,
         rng: &mut R,
     ) -> Result<Transcript> {
-        self.0.create_share(rng)
+        match self.0.share(rng) {
+            Ok(Message::Deal(transcript)) => Ok(transcript),
+            Err(e) => Err(e),
+            _ => Err(Error::InvalidDkgStateToDeal),
+        }
     }
 
     pub fn aggregate_transcripts(
@@ -435,7 +439,7 @@ mod test_ferveo_api {
         let messages: Vec<_> = validators
             .iter()
             .map(|sender| {
-                let dkg = Dkg::new(
+                let mut dkg = Dkg::new(
                     tau,
                     shares_num,
                     security_threshold,
diff --git a/ferveo/src/bindings_python.rs b/ferveo/src/bindings_python.rs
index 69fe1b8d..c35dc291 100644
--- a/ferveo/src/bindings_python.rs
+++ b/ferveo/src/bindings_python.rs
@@ -491,7 +491,7 @@ impl Dkg {
         DkgPublicKey(self.0.public_key())
     }
 
-    pub fn generate_transcript(&self) -> PyResult<Transcript> {
+    pub fn generate_transcript(&mut self) -> PyResult<Transcript> {
         let rng = &mut thread_rng();
         let transcript = self
             .0
@@ -776,7 +776,7 @@ mod test_ferveo_python {
             .iter()
             .cloned()
             .map(|sender| {
-                let dkg = Dkg::new(
+                let mut dkg = Dkg::new(
                     tau,
                     shares_num,
                     security_threshold,
diff --git a/ferveo/src/bindings_wasm.rs b/ferveo/src/bindings_wasm.rs
index 1c0d2a0b..5a23909a 100644
--- a/ferveo/src/bindings_wasm.rs
+++ b/ferveo/src/bindings_wasm.rs
@@ -366,7 +366,7 @@ impl Dkg {
     }
 
     #[wasm_bindgen(js_name = "generateTranscript")]
-    pub fn generate_transcript(&self) -> JsResult<Transcript> {
+    pub fn generate_transcript(&mut self) -> JsResult<Transcript> {
         let rng = &mut thread_rng();
         let transcript = self.0.generate_transcript(rng).map_err(map_js_err)?;
         Ok(Transcript(transcript))
diff --git a/ferveo/src/dkg.rs b/ferveo/src/dkg.rs
index f3c0dc15..10f73d7e 100644
--- a/ferveo/src/dkg.rs
+++ b/ferveo/src/dkg.rs
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 
 use ark_ec::{pairing::Pairing, AffineRepr, CurveGroup};
 use ark_poly::EvaluationDomain;
+use ark_std::UniformRand;
 use ferveo_common::PublicKey;
 use measure_time::print_time;
 use rand::RngCore;
@@ -159,24 +160,19 @@ impl<E: Pairing> PubliclyVerifiableDkg<E> {
     /// Returns a PVSS dealing message to post on-chain
     pub fn share<R: RngCore>(&mut self, rng: &mut R) -> Result<Message<E>> {
         print_time!("PVSS Sharing");
-        let vss = self.create_share(rng)?;
         match self.state {
             DkgState::Sharing { .. } | DkgState::Dealt => {
+                let vss = PubliclyVerifiableSS::<E>::new(
+                    &E::ScalarField::rand(rng),
+                    self,
+                    rng,
+                )?;
                 Ok(Message::Deal(vss))
             }
             _ => Err(Error::InvalidDkgStateToDeal),
         }
     }
 
-    // TODO: Make private, use `share` instead. Currently used only in bindings
-    pub fn create_share<R: RngCore>(
-        &self,
-        rng: &mut R,
-    ) -> Result<PubliclyVerifiableSS<E>> {
-        use ark_std::UniformRand;
-        PubliclyVerifiableSS::<E>::new(&E::ScalarField::rand(rng), self, rng)
-    }
-
     /// Aggregate all received PVSS messages into a single message, prepared to post on-chain
     pub fn aggregate(&self) -> Result<Message<E>> {
         match self.state {