Enforce FROST StateMachine progression via the type system

A comment on the matter was made in 
https://github.com/serai-dex/serai/issues/12. While I do believe the API 
is slightly worse, I appreciate the explicitness.

crypto/frost/src/sign.rs

@@ -236,31 +236,21 @@ fn complete<C: Curve, A: Algorithm<C>>(
   )
 }
 
-/// State of a Sign machine
-#[derive(Clone, Copy, PartialEq, Eq, Debug)]
-pub enum State {
-  Fresh,
-  Preprocessed,
-  Signed,
-  Complete,
-}
-
-impl fmt::Display for State {
-  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-    write!(f, "{:?}", self)
-  }
-}
-
-pub trait StateMachine {
+pub trait PreprocessMachine {
   type Signature: Clone + PartialEq + fmt::Debug;
+  type SignMachine: SignMachine<Self::Signature>;
 
   /// Perform the preprocessing round required in order to sign
   /// Returns a byte vector which must be transmitted to all parties selected for this signing
   /// process, over an authenticated channel
   fn preprocess<R: RngCore + CryptoRng>(
-    &mut self,
+    self,
     rng: &mut R
-  ) -> Result<Vec<u8>, FrostError>;
+  ) -> (Self::SignMachine, Vec<u8>);
+}
+
+pub trait SignMachine<S> {
+  type SignatureMachine: SignatureMachine<S>;
 
   /// Sign a message
   /// Takes in the participant's commitments, which are expected to be in a Vec where participant
@@ -268,29 +258,33 @@ pub trait StateMachine {
   /// index i which is locally handled. Returns a byte vector representing a share of the signature
   /// for every other participant to receive, over an authenticated channel
   fn sign(
-    &mut self,
+    self,
     commitments: HashMap<u16, Vec<u8>>,
     msg: &[u8],
-  ) -> Result<Vec<u8>, FrostError>;
+  ) -> Result<(Self::SignatureMachine, Vec<u8>), FrostError>;
+}
 
+pub trait SignatureMachine<S> {
   /// Complete signing
   /// Takes in everyone elses' shares submitted to us as a Vec, expecting participant index =
   /// Vec index with None at index 0 and index i. Returns a byte vector representing the serialized
   /// signature
-  fn complete(&mut self, shares: HashMap<u16, Vec<u8>>) -> Result<Self::Signature, FrostError>;
-
-  fn multisig_params(&self) -> MultisigParams;
-
-  fn state(&self) -> State;
+  fn complete(self, shares: HashMap<u16, Vec<u8>>) -> Result<S, FrostError>;
 }
 
 /// State machine which manages signing for an arbitrary signature algorithm
-#[allow(non_snake_case)]
 pub struct AlgorithmMachine<C: Curve, A: Algorithm<C>> {
+  params: Params<C, A>
+}
+
+pub struct AlgorithmSignMachine<C: Curve, A: Algorithm<C>> {
   params: Params<C, A>,
-  state: State,
-  preprocess: Option<PreprocessPackage<C>>,
-  sign: Option<Package<C>>,
+  preprocess: PreprocessPackage<C>,
+}
+
+pub struct AlgorithmSignatureMachine<C: Curve, A: Algorithm<C>> {
+  params: Params<C, A>,
+  sign: Package<C>,
 }
 
 impl<C: Curve, A: Algorithm<C>> AlgorithmMachine<C, A> {
@@ -300,85 +294,52 @@ impl<C: Curve, A: Algorithm<C>> AlgorithmMachine<C, A> {
     keys: Arc<MultisigKeys<C>>,
     included: &[u16],
   ) -> Result<AlgorithmMachine<C, A>, FrostError> {
-    Ok(
-      AlgorithmMachine {
-        params: Params::new(algorithm, keys, included)?,
-        state: State::Fresh,
-        preprocess: None,
-        sign: None,
-      }
-    )
+    Ok(AlgorithmMachine { params: Params::new(algorithm, keys, included)? })
   }
 
-  pub(crate) fn unsafe_override_preprocess(&mut self, preprocess: PreprocessPackage<C>) {
-    if self.state != State::Fresh {
-      // This would be unacceptable, yet this is pub(crate) and explicitly labelled unsafe
-      // It's solely used in a testing environment, which is how it's justified
-      Err::<(), _>(FrostError::InvalidSignTransition(State::Fresh, self.state)).unwrap();
-    }
-    self.preprocess = Some(preprocess);
-    self.state = State::Preprocessed;
+  pub(crate) fn unsafe_override_preprocess(
+    self,
+    preprocess: PreprocessPackage<C>
+  ) -> (AlgorithmSignMachine<C, A>, Vec<u8>) {
+    let serialized = preprocess.serialized.clone();
+    (AlgorithmSignMachine { params: self.params, preprocess }, serialized)
   }
 }
 
-impl<C: Curve, A: Algorithm<C>> StateMachine for AlgorithmMachine<C, A> {
+impl<C: Curve, A: Algorithm<C>> PreprocessMachine for AlgorithmMachine<C, A> {
   type Signature = A::Signature;
+  type SignMachine = AlgorithmSignMachine<C, A>;
 
   fn preprocess<R: RngCore + CryptoRng>(
-    &mut self,
+    self,
     rng: &mut R
-  ) -> Result<Vec<u8>, FrostError> {
-    if self.state != State::Fresh {
-      Err(FrostError::InvalidSignTransition(State::Fresh, self.state))?;
-    }
-    let preprocess = preprocess::<R, C, A>(rng, &mut self.params);
+  ) -> (Self::SignMachine, Vec<u8>) {
+    let mut params = self.params;
+    let preprocess = preprocess::<R, C, A>(rng, &mut params);
     let serialized = preprocess.serialized.clone();
-    self.preprocess = Some(preprocess);
-    self.state = State::Preprocessed;
-    Ok(serialized)
-  }
-
-  fn sign(
-    &mut self,
-    commitments: HashMap<u16, Vec<u8>>,
-    msg: &[u8],
-  ) -> Result<Vec<u8>, FrostError> {
-    if self.state != State::Preprocessed {
-      Err(FrostError::InvalidSignTransition(State::Preprocessed, self.state))?;
-    }
-
-    let (sign, serialized) = sign_with_share(
-      &mut self.params,
-      self.preprocess.take().unwrap(),
-      commitments,
-      msg,
-    )?;
-
-    self.sign = Some(sign);
-    self.state = State::Signed;
-    Ok(serialized)
-  }
-
-  fn complete(&mut self, shares: HashMap<u16, Vec<u8>>) -> Result<A::Signature, FrostError> {
-    if self.state != State::Signed {
-      Err(FrostError::InvalidSignTransition(State::Signed, self.state))?;
-    }
-
-    let signature = complete(
-      &self.params,
-      self.sign.take().unwrap(),
-      shares,
-    )?;
-
-    self.state = State::Complete;
-    Ok(signature)
-  }
-
-  fn multisig_params(&self) -> MultisigParams {
-    self.params.multisig_params().clone()
-  }
-
-  fn state(&self) -> State {
-    self.state
+    (AlgorithmSignMachine { params, preprocess }, serialized)
+  }
+}
+
+impl<C: Curve, A: Algorithm<C>> SignMachine<A::Signature> for AlgorithmSignMachine<C, A> {
+  type SignatureMachine = AlgorithmSignatureMachine<C, A>;
+
+  fn sign(
+    self,
+    commitments: HashMap<u16, Vec<u8>>,
+    msg: &[u8]
+  ) -> Result<(Self::SignatureMachine, Vec<u8>), FrostError> {
+    let mut params = self.params;
+    let (sign, serialized) = sign_with_share(&mut params, self.preprocess, commitments, msg)?;
+    Ok((AlgorithmSignatureMachine { params, sign }, serialized))
+  }
+}
+
+impl<
+  C: Curve,
+  A: Algorithm<C>
+> SignatureMachine<A::Signature> for AlgorithmSignatureMachine<C, A> {
+  fn complete(self, shares: HashMap<u16, Vec<u8>>) -> Result<A::Signature, FrostError> {
+    complete(&self.params, self.sign, shares)
   }
 }