Move FMD component to separate module

This fixes the WASM build.
penumbra-zone · May 9, 2024 · cfd3a35 · cfd3a35
1 parent 7244359
commit cfd3a35
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 90 deletions.
diff --git a/crates/core/component/shielded-pool/src/component.rs b/crates/core/component/shielded-pool/src/component.rs
@@ -2,13 +2,15 @@
 
 mod action_handler;
 mod assets;
+mod fmd;
 mod metrics;
 mod note_manager;
 mod shielded_pool;
 mod transfer;
 
 pub use self::metrics::register_metrics;
 pub use assets::{AssetRegistry, AssetRegistryRead};
+pub use fmd::ClueManager;
 pub use note_manager::NoteManager;
 pub use shielded_pool::{ShieldedPool, StateReadExt, StateWriteExt};
 pub use transfer::Ics20Transfer;

diff --git a/crates/core/component/shielded-pool/src/component/fmd.rs b/crates/core/component/shielded-pool/src/component/fmd.rs
@@ -0,0 +1,92 @@
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use cnidarium::{StateRead, StateWrite};
+use decaf377_fmd::Clue;
+use penumbra_proto::{
+    core::component::shielded_pool::v1::{self as pb},
+    StateWriteProto,
+};
+use penumbra_txhash::TransactionId;
+
+use crate::fmd::state_key;
+
+#[async_trait]
+trait ClueWriteExt: StateWrite {
+    fn put_current_clue_count(&mut self, count: u64) {
+        self.put_raw(
+            state_key::clue_count::current().to_string(),
+            count.to_be_bytes().to_vec(),
+        )
+    }
+
+    fn put_previous_clue_count(&mut self, count: u64) {
+        self.put_raw(
+            state_key::clue_count::previous().to_string(),
+            count.to_be_bytes().to_vec(),
+        )
+    }
+}
+
+impl<T: StateWrite + ?Sized> ClueWriteExt for T {}
+
+#[async_trait]
+trait ClueReadExt: StateRead {
+    async fn get_current_clue_count(&self) -> Result<u64> {
+        Ok(u64::from_be_bytes(
+            self.get_raw(state_key::clue_count::current())
+                .await?
+                .ok_or(anyhow!("no current clue count"))?
+                .as_slice()
+                .try_into()?,
+        ))
+    }
+
+    async fn get_previous_clue_count(&self) -> Result<u64> {
+        Ok(u64::from_be_bytes(
+            self.get_raw(state_key::clue_count::previous())
+                .await?
+                .ok_or(anyhow!("no current clue count"))?
+                .as_slice()
+                .try_into()?,
+        ))
+    }
+}
+
+impl<T: StateRead + ?Sized> ClueReadExt for T {}
+
+#[async_trait]
+pub trait ClueManager: StateRead + StateWrite {
+    async fn record_clue(&mut self, clue: Clue, tx: TransactionId) -> Result<()> {
+        // Update count
+        {
+            let count = self.get_current_clue_count().await?;
+            self.put_current_clue_count(count.saturating_add(1));
+        }
+        self.record_proto(pb::EventClue {
+            clue: Some(clue.into()),
+            tx: Some(tx.into()),
+        });
+        Ok(())
+    }
+}
+
+impl<T: StateRead + StateWrite> ClueManager for T {}
+
+#[async_trait]
+pub(crate) trait ClueManagerInternal: ClueManager {
+    fn init(&mut self) {
+        self.put_current_clue_count(0);
+        self.put_previous_clue_count(0);
+    }
+
+    /// Flush the clue counts, returning the previous and current counts
+    async fn flush_clue_count(&mut self) -> Result<(u64, u64)> {
+        let previous = self.get_previous_clue_count().await?;
+        let current = self.get_current_clue_count().await?;
+        self.put_previous_clue_count(current);
+        self.put_current_clue_count(0);
+        Ok((previous, current))
+    }
+}
+
+impl<T: ClueManager> ClueManagerInternal for T {}
diff --git a/crates/core/component/shielded-pool/src/component/shielded_pool.rs b/crates/core/component/shielded-pool/src/component/shielded_pool.rs
@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
-use crate::fmd::{should_update_fmd_params, ClueManagerInternal as _};
+use super::fmd::ClueManagerInternal as _;
+use crate::fmd::should_update_fmd_params;
 use crate::params::ShieldedPoolParameters;
 use crate::{fmd, genesis, state_key};
 use anyhow::anyhow;

diff --git a/crates/core/component/shielded-pool/src/fmd.rs b/crates/core/component/shielded-pool/src/fmd.rs
@@ -1,12 +1,9 @@
 use anyhow::{anyhow, Result};
-use async_trait::async_trait;
-use cnidarium::{StateRead, StateWrite};
-use decaf377_fmd::{Clue, Precision};
+use decaf377_fmd::Precision;
 use penumbra_proto::{
     core::component::shielded_pool::v1::{self as pb},
-    DomainType, StateWriteProto,
+    DomainType,
 };
-use penumbra_txhash::TransactionId;
 use serde::{Deserialize, Serialize};
 
 pub mod state_key;
@@ -107,7 +104,12 @@ impl Default for MetaParameters {
 }
 
 impl MetaParameters {
-    pub fn updated_fmd_params(&self, _old: &Parameters, height: u64, _clue_count_delta: (u64, u64)) -> Parameters {
+    pub fn updated_fmd_params(
+        &self,
+        _old: &Parameters,
+        height: u64,
+        _clue_count_delta: (u64, u64),
+    ) -> Parameters {
         match *self {
             MetaParameters::Fixed(precision) => Parameters {
                 precision,
@@ -116,84 +118,3 @@ impl MetaParameters {
         }
     }
 }
-
-#[async_trait]
-trait ClueWriteExt: StateWrite {
-    fn put_current_clue_count(&mut self, count: u64) {
-        self.put_raw(
-            state_key::clue_count::current().to_string(),
-            count.to_be_bytes().to_vec(),
-        )
-    }
-
-    fn put_previous_clue_count(&mut self, count: u64) {
-        self.put_raw(
-            state_key::clue_count::previous().to_string(),
-            count.to_be_bytes().to_vec(),
-        )
-    }
-}
-
-impl<T: StateWrite + ?Sized> ClueWriteExt for T {}
-
-#[async_trait]
-trait ClueReadExt: StateRead {
-    async fn get_current_clue_count(&self) -> Result<u64> {
-        Ok(u64::from_be_bytes(
-            self.get_raw(state_key::clue_count::current())
-                .await?
-                .ok_or(anyhow!("no current clue count"))?
-                .as_slice()
-                .try_into()?,
-        ))
-    }
-
-    async fn get_previous_clue_count(&self) -> Result<u64> {
-        Ok(u64::from_be_bytes(
-            self.get_raw(state_key::clue_count::previous())
-                .await?
-                .ok_or(anyhow!("no current clue count"))?
-                .as_slice()
-                .try_into()?,
-        ))
-    }
-}
-
-impl<T: StateRead + ?Sized> ClueReadExt for T {}
-
-#[async_trait]
-pub trait ClueManager: StateRead + StateWrite {
-    async fn record_clue(&mut self, clue: Clue, tx: TransactionId) -> Result<()> {
-        // Update count
-        {
-            let count = self.get_current_clue_count().await?;
-            self.put_current_clue_count(count.saturating_add(1));
-        }
-        self.record_proto(pb::EventClue {
-            clue: Some(clue.into()),
-            tx: Some(tx.into())
-        });
-        Ok(())
-    }
-}
-
-impl<T: StateRead + StateWrite> ClueManager for T {}
-
-#[async_trait]
-pub(crate) trait ClueManagerInternal: ClueManager {
-    fn init(&mut self) {
-        self.put_current_clue_count(0);
-        self.put_previous_clue_count(0);
-    }
-
-    /// Flush the clue counts, returning the previous and current counts
-    async fn flush_clue_count(&mut self) -> Result<(u64, u64)> {
-        let previous = self.get_previous_clue_count().await?;
-        let current = self.get_current_clue_count().await?;
-        self.put_previous_clue_count(current);
-        self.put_current_clue_count(0);
-        Ok((previous, current))
-    }
-}
-
-impl<T: ClueManager> ClueManagerInternal for T {}
diff --git a/crates/core/component/shielded-pool/src/fmd/state_key.rs b/crates/core/component/shielded-pool/src/fmd/state_key.rs
@@ -10,7 +10,7 @@ pub mod parameters {
     }
 }
 
-pub(super) mod clue_count {
+pub(crate) mod clue_count {
     pub fn current() -> &'static str {
         "shielded_pool/fmd_clue_count/current"
     }

diff --git a/crates/core/component/shielded-pool/src/lib.rs b/crates/core/component/shielded-pool/src/lib.rs
@@ -26,7 +26,6 @@ pub mod output;
 pub mod spend;
 
 pub use convert::{ConvertCircuit, ConvertProof, ConvertProofPrivate, ConvertProofPublic};
-pub use fmd::ClueManager;
 pub use nullifier_derivation::{
     NullifierDerivationCircuit, NullifierDerivationProof, NullifierDerivationProofPrivate,
     NullifierDerivationProofPublic,