refactor(wasm): modularize wasm crate logic (#3432)

* update wasm npm publish script

* add LTO flag

* add bin folder to npm package

* migrated 'build_action' to dedicated file

* clippy and fmt

* simplified 'build_action' wasm method

* removed unused imports to pass CI

* fmt transaction package

* addressing gabe's comments

* reverting NPM package changes

crates/core/transaction/src/lib.rs

@@ -37,6 +37,7 @@ pub use effect_hash::EffectingData;
 pub use error::Error;
 pub use id::Id;
 pub use is_action::IsAction;
+pub use plan::ActionPlan;
 pub use transaction::{Transaction, TransactionBody, TransactionParameters};
 pub use view::{ActionView, MemoPlaintextView, MemoView, TransactionPerspective, TransactionView};
 pub use witness_data::WitnessData;

crates/core/transaction/src/plan/action.rs

@@ -91,14 +91,14 @@ impl ActionPlan {
     /// This method is useful for controlling how a transaction's actions are
     /// built (e.g., building them in parallel, or via Web Workers).
     pub fn build_unauth(
-        &self,
+        action_plan: ActionPlan,
         fvk: &FullViewingKey,
         witness_data: &WitnessData,
         memo_key: Option<PayloadKey>,
     ) -> Result<Action> {
         use ActionPlan::*;
 
-        Ok(match self {
+        Ok(match action_plan {
             Spend(spend_plan) => {
                 let note_commitment = spend_plan.note.commit();
                 let auth_path = witness_data
@@ -130,7 +130,7 @@ impl ActionPlan {
                     .get(&note_commitment)
                     .context(format!("could not get proof for {note_commitment:?}"))?;
 
-                Action::SwapClaim(swap_claim_plan.swap_claim(&fvk, auth_path))
+                Action::SwapClaim(swap_claim_plan.swap_claim(fvk, auth_path))
             }
             Delegate(plan) => Action::Delegate(plan.clone()),
             Undelegate(plan) => Action::Undelegate(plan.clone()),

crates/core/transaction/src/plan/build.rs

@@ -8,6 +8,7 @@ use rand_core::{CryptoRng, RngCore};
 use std::fmt::Debug;
 
 use super::TransactionPlan;
+use crate::ActionPlan;
 use crate::{
     action::Action,
     transaction::{DetectionData, TransactionParameters},
@@ -143,7 +144,12 @@ impl TransactionPlan {
             .actions
             .iter()
             .map(|action_plan| {
-                action_plan.build_unauth(full_viewing_key, witness_data, self.memo_key())
+                ActionPlan::build_unauth(
+                    action_plan.clone(),
+                    full_viewing_key,
+                    witness_data,
+                    self.memo_key(),
+                )
             })
             .collect::<Result<Vec<_>>>()?;
 
@@ -181,7 +187,7 @@ impl TransactionPlan {
                 let witness_data2 = witness_data.clone(); // Arc
                 let memo_key2 = self.memo_key();
                 tokio::task::spawn_blocking(move || {
-                    action_plan.build_unauth(&fvk2, &*witness_data2, memo_key2)
+                    ActionPlan::build_unauth(action_plan, &fvk2, &*witness_data2, memo_key2)
                 })
             })
             .collect::<Vec<_>>();

crates/core/transaction/src/witness_data.rs

@@ -3,8 +3,10 @@ use std::collections::BTreeMap;
 use penumbra_proto::{core::transaction::v1alpha1 as pb, DomainType, TypeUrl};
 use penumbra_shielded_pool::note;
 use penumbra_tct as tct;
+use serde::{Deserialize, Serialize};
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(try_from = "pb::WitnessData", into = "pb::WitnessData")]
 pub struct WitnessData {
     pub anchor: tct::Root,
     pub state_commitment_proofs: BTreeMap<note::StateCommitment, tct::Proof>,

crates/wasm/Cargo.toml

@@ -49,4 +49,7 @@ web-sys                  = { version = "0.3.64", features = ["console"] }
 
 [dev-dependencies]
 wasm-bindgen-test        = "0.3.37"
-serde_json               = "1.0.107"
+serde_json               = "1.0.107"
+
+[profile.release]
+lto = true

crates/wasm/publish/package.json

@@ -12,9 +12,9 @@
     "wasm-pack": "^0.12.1"
   },
   "devDependencies": {
-    "@types/node": "^20.5.6",
+    "@types/node": "^20.9.4",
+    "prettier": "^3.0.2",
     "tsx": "^3.12.7",
-    "typescript": "^5.2.2",
-    "prettier": "^3.0.2"
+    "typescript": "^5.2.2"
   }
 }

crates/wasm/publish/run.ts

@@ -30,4 +30,4 @@ TARGETS.forEach(target => {
 
   // Change working directory back to parent
   process.chdir('..');
-});
+});

crates/wasm/src/build.rs

@@ -0,0 +1,45 @@
+use crate::error::WasmResult;
+use std::str::FromStr;
+use wasm_bindgen::prelude::wasm_bindgen;
+use wasm_bindgen::JsValue;
+
+use crate::utils;
+use penumbra_keys::FullViewingKey;
+use penumbra_transaction::{
+    plan::{ActionPlan, TransactionPlan},
+    WitnessData,
+};
+
+/// Builds a planned [`Action`] specified by
+/// the [`ActionPlan`] in a [`TransactionPlan`].
+/// Arguments:
+///     transaction_plan: `TransactionPlan`
+///     action_plan: `ActionPlan`
+///     full_viewing_key: `bech32m String`,
+///     witness_data: `WitnessData``
+/// Returns: `Action`
+#[wasm_bindgen]
+pub fn build_action(
+    transaction_plan: JsValue,
+    action_plan: JsValue,
+    full_viewing_key: &str,
+    witness_data: JsValue,
+) -> WasmResult<JsValue> {
+    utils::set_panic_hook();
+
+    let transaction_plan: TransactionPlan =
+        serde_wasm_bindgen::from_value(transaction_plan.clone())?;
+
+    let witness: WitnessData = serde_wasm_bindgen::from_value(witness_data)?;
+
+    let action_plan: ActionPlan = serde_wasm_bindgen::from_value(action_plan)?;
+
+    let full_viewing_key: FullViewingKey = FullViewingKey::from_str(full_viewing_key)?;
+
+    let memo_key = transaction_plan.memo_plan.map(|memo_plan| memo_plan.key);
+
+    let action = ActionPlan::build_unauth(action_plan, &full_viewing_key, &witness, memo_key)?;
+
+    let result = serde_wasm_bindgen::to_value(&action)?;
+    Ok(result)
+}

crates/wasm/src/lib.rs

@@ -1,6 +1,7 @@
 #![allow(dead_code)]
 extern crate core;
 
+pub mod build;
 pub mod error;
 pub mod keys;
 pub mod note_record;

crates/wasm/src/wasm_planner.rs

@@ -1,6 +1,4 @@
-use std::str::FromStr;
-
-use anyhow::{anyhow, Context};
+use anyhow::anyhow;
 use ark_ff::UniformRand;
 use decaf377::Fq;
 use rand_core::OsRng;
@@ -9,20 +7,17 @@ use wasm_bindgen::JsValue;
 
 use penumbra_chain::params::{ChainParameters, FmdParameters};
 use penumbra_dex::swap_claim::SwapClaimPlan;
-use penumbra_keys::FullViewingKey;
 use penumbra_proto::{
     core::{
         asset::v1alpha1::{DenomMetadata, Value},
         component::fee::v1alpha1::{Fee, GasPrices},
         component::ibc::v1alpha1::Ics20Withdrawal,
         keys::v1alpha1::{Address, AddressIndex},
-        transaction::v1alpha1 as pb,
         transaction::v1alpha1::MemoPlaintext,
     },
     crypto::tct::v1alpha1::StateCommitment,
     DomainType,
 };
-use penumbra_transaction::{action::Action, plan::ActionPlan, plan::TransactionPlan, WitnessData};
 
 use crate::error::WasmResult;
 use crate::planner::Planner;
@@ -65,112 +60,6 @@ impl WasmPlanner {
         Ok(planner)
     }
 
-    /// Builds a planned [`Action`] specified by
-    /// the [`ActionPlan`] in a [`TransactionPlan`].
-    /// Arguments:
-    ///     &self: `WasmPlanner`
-    ///     transaction_plan: `TransactionPlan`
-    ///     action_plan: `ActionPlan`
-    ///     full_viewing_key: `bech32m String`,
-    ///     witness_data: `WitnessData``
-    /// Returns: `Action`
-    #[wasm_bindgen]
-    pub fn build_action(
-        &self,
-        transaction_plan: JsValue,
-        action_plan: JsValue,
-        full_viewing_key: &str,
-        witness_data: JsValue,
-    ) -> WasmResult<JsValue> {
-        utils::set_panic_hook();
-
-        let transaction_plan: TransactionPlan =
-            serde_wasm_bindgen::from_value(transaction_plan.clone())?;
-
-        let witness_data_proto: pb::WitnessData = serde_wasm_bindgen::from_value(witness_data)?;
-        let witness_data: WitnessData = witness_data_proto.try_into()?;
-
-        let action_plan: ActionPlan = serde_wasm_bindgen::from_value(action_plan)?;
-
-        let full_viewing_key: FullViewingKey = FullViewingKey::from_str(full_viewing_key)?;
-
-        let memo_key = transaction_plan.memo_plan.map(|memo_plan| memo_plan.key);
-
-        let action = match action_plan {
-            ActionPlan::Spend(spend_plan) => {
-                let spend = ActionPlan::Spend(spend_plan);
-                Some(spend.build_unauth(&full_viewing_key, &witness_data, memo_key)?)
-            }
-            ActionPlan::Output(output_plan) => {
-                let output = ActionPlan::Output(output_plan);
-                Some(output.build_unauth(&full_viewing_key, &witness_data, memo_key)?)
-            }
-
-            // TODO: Other action variants besides 'Spend' and 'Output' still require testing.
-            ActionPlan::Swap(swap_plan) => {
-                let swap = ActionPlan::Swap(swap_plan);
-                Some(swap.build_unauth(&full_viewing_key, &witness_data, memo_key)?)
-            }
-            ActionPlan::SwapClaim(swap_claim_plan) => {
-                let swap_claim = ActionPlan::SwapClaim(swap_claim_plan);
-                Some(swap_claim.build_unauth(&full_viewing_key, &witness_data, memo_key)?)
-            }
-            ActionPlan::Delegate(delegation) => Some(Action::Delegate(delegation)),
-            ActionPlan::Undelegate(undelegation) => Some(Action::Undelegate(undelegation)),
-            ActionPlan::UndelegateClaim(undelegate_claim) => {
-                let undelegate_claim = undelegate_claim.undelegate_claim();
-                Some(Action::UndelegateClaim(undelegate_claim))
-            }
-            ActionPlan::ProposalSubmit(proposal_submit) => {
-                Some(Action::ProposalSubmit(proposal_submit))
-            }
-            ActionPlan::ProposalWithdraw(proposal_withdraw) => {
-                Some(Action::ProposalWithdraw(proposal_withdraw))
-            }
-            ActionPlan::ValidatorVote(validator_vote) => {
-                Some(Action::ValidatorVote(validator_vote))
-            }
-            ActionPlan::DelegatorVote(delegator_vote) => {
-                let note_commitment = delegator_vote.staked_note.commit();
-                let auth_path = witness_data
-                    .state_commitment_proofs
-                    .get(&note_commitment)
-                    .context(format!("could not get proof for {note_commitment:?}"))?;
-
-                Some(Action::DelegatorVote(delegator_vote.delegator_vote(
-                    &full_viewing_key,
-                    [0; 64].into(),
-                    auth_path.clone(),
-                )))
-            }
-            ActionPlan::ProposalDepositClaim(proposal_deposit_claim) => {
-                Some(Action::ProposalDepositClaim(proposal_deposit_claim))
-            }
-            ActionPlan::ValidatorDefinition(validator_definition) => {
-                Some(Action::ValidatorDefinition(validator_definition))
-            }
-            ActionPlan::IbcAction(ibc_action) => Some(Action::IbcRelay(ibc_action)),
-            ActionPlan::DaoSpend(dao_spend) => Some(Action::DaoSpend(dao_spend)),
-            ActionPlan::DaoOutput(dao_output) => Some(Action::DaoOutput(dao_output)),
-            ActionPlan::DaoDeposit(dao_deposit) => Some(Action::DaoDeposit(dao_deposit)),
-            ActionPlan::PositionOpen(position_open) => Some(Action::PositionOpen(position_open)),
-            ActionPlan::PositionClose(position_close) => {
-                Some(Action::PositionClose(position_close))
-            }
-            ActionPlan::PositionWithdraw(position_withdrawn) => Some(Action::PositionWithdraw(
-                position_withdrawn.position_withdraw(),
-            )),
-            ActionPlan::Withdrawal(ics20_withdrawal) => {
-                Some(Action::Ics20Withdrawal(ics20_withdrawal))
-            }
-            // TODO: Should we handle `PositionRewardClaim`?
-            _ => None,
-        };
-
-        let action_result_proto = serde_wasm_bindgen::to_value(&Some(action))?;
-        Ok(action_result_proto)
-    }
-
     /// Public getter for the 'storage' field
     pub fn get_storage(&self) -> *const IndexedDBStorage {
         &self.storage
@@ -179,7 +68,6 @@ impl WasmPlanner {
     /// Add expiry height to plan
     /// Arguments:
     ///     expiry_height: `u64`
-    #[wasm_bindgen]
     pub fn expiry_height(&mut self, expiry_height: u64) -> WasmResult<()> {
         utils::set_panic_hook();
 
@@ -237,7 +125,6 @@ impl WasmPlanner {
     /// Add swap claim to plan
     /// Arguments:
     ///     swap_commitment: `StateCommitment`
-    #[wasm_bindgen]
     pub async fn swap_claim(&mut self, swap_commitment: JsValue) -> WasmResult<()> {
         utils::set_panic_hook();
 

crates/wasm/tests/test_build.rs

@@ -27,6 +27,7 @@ mod tests {
         Action,
     };
     use penumbra_wasm::{
+        build::build_action,
         error::WasmError,
         keys::load_proving_key,
         storage::IndexedDBStorage,
@@ -451,28 +452,26 @@ mod tests {
         for i in transaction_plan_conv.actions.clone() {
             if let ActionPlan::Spend(ref _spend_plan) = i {
                 let action_deserialize = serde_wasm_bindgen::to_value(&i).unwrap();
-                let action = wasm_planner
-                    .build_action(
-                        transaction_plan.clone(),
-                        action_deserialize,
-                        full_viewing_key,
-                        witness_data.as_ref().unwrap().clone(),
-                    )
-                    .unwrap();
+                let action = build_action(
+                    transaction_plan.clone(),
+                    action_deserialize,
+                    full_viewing_key,
+                    witness_data.as_ref().unwrap().clone(),
+                )
+                .unwrap();
                 let action_serialize: Action =
                     serde_wasm_bindgen::from_value(action.clone()).unwrap();
                 actions.push(action_serialize);
             }
             if let ActionPlan::Output(ref _output_plan) = i {
                 let action_deserialize = serde_wasm_bindgen::to_value(&i).unwrap();
-                let action = wasm_planner
-                    .build_action(
-                        transaction_plan.clone(),
-                        action_deserialize,
-                        full_viewing_key,
-                        witness_data.as_ref().unwrap().clone(),
-                    )
-                    .unwrap();
+                let action = build_action(
+                    transaction_plan.clone(),
+                    action_deserialize,
+                    full_viewing_key,
+                    witness_data.as_ref().unwrap().clone(),
+                )
+                .unwrap();
                 let action_serialize: Action =
                     serde_wasm_bindgen::from_value(action.clone()).unwrap();
                 actions.push(action_serialize);