perf: Make PublicKey decoding lazy inside WASM

Signed-off-by: Dmitry Murzin <[email protected]>

crypto/src/lib.rs

@@ -38,6 +38,8 @@ pub use hash::*;
 use iroha_macro::ffi_impl_opaque;
 use iroha_primitives::const_vec::ConstVec;
 use iroha_schema::{Declaration, IntoSchema, MetaMap, Metadata, NamedFieldsMeta, TypeId};
+#[cfg(target_family = "wasm")]
+use lazy::PublicKeyLazy;
 pub use merkle::MerkleTree;
 #[cfg(not(feature = "ffi_import"))]
 use parity_scale_codec::{Decode, Encode};
@@ -248,7 +250,7 @@ impl From<PrivateKey> for KeyPair {
 impl From<(ed25519::PublicKey, ed25519::PrivateKey)> for KeyPair {
     fn from((public_key, private_key): (ed25519::PublicKey, ed25519::PrivateKey)) -> Self {
         Self {
-            public_key: PublicKey(Box::new(PublicKeyInner::Ed25519(public_key))),
+            public_key: PublicKey::new(PublicKeyInner::Ed25519(public_key)),
             private_key: PrivateKey(Box::new(Secret::new(PrivateKeyInner::Ed25519(private_key)))),
         }
     }
@@ -257,7 +259,7 @@ impl From<(ed25519::PublicKey, ed25519::PrivateKey)> for KeyPair {
 impl From<(secp256k1::PublicKey, secp256k1::PrivateKey)> for KeyPair {
     fn from((public_key, private_key): (secp256k1::PublicKey, secp256k1::PrivateKey)) -> Self {
         Self {
-            public_key: PublicKey(Box::new(PublicKeyInner::Secp256k1(public_key))),
+            public_key: PublicKey::new(PublicKeyInner::Secp256k1(public_key)),
             private_key: PrivateKey(Box::new(Secret::new(PrivateKeyInner::Secp256k1(
                 private_key,
             )))),
@@ -270,7 +272,7 @@ impl From<(bls::BlsNormalPublicKey, bls::BlsNormalPrivateKey)> for KeyPair {
         (public_key, private_key): (bls::BlsNormalPublicKey, bls::BlsNormalPrivateKey),
     ) -> Self {
         Self {
-            public_key: PublicKey(Box::new(PublicKeyInner::BlsNormal(public_key))),
+            public_key: PublicKey::new(PublicKeyInner::BlsNormal(public_key)),
             private_key: PrivateKey(Box::new(Secret::new(PrivateKeyInner::BlsNormal(
                 private_key,
             )))),
@@ -281,7 +283,7 @@ impl From<(bls::BlsNormalPublicKey, bls::BlsNormalPrivateKey)> for KeyPair {
 impl From<(bls::BlsSmallPublicKey, bls::BlsSmallPrivateKey)> for KeyPair {
     fn from((public_key, private_key): (bls::BlsSmallPublicKey, bls::BlsSmallPrivateKey)) -> Self {
         Self {
-            public_key: PublicKey(Box::new(PublicKeyInner::BlsSmall(public_key))),
+            public_key: PublicKey::new(PublicKeyInner::BlsSmall(public_key)),
             private_key: PrivateKey(Box::new(Secret::new(PrivateKeyInner::BlsSmall(
                 private_key,
             )))),
@@ -346,7 +348,7 @@ impl FromStr for PublicKeyInner {
         let bytes = hex_decode(key)?;
 
         let (algorithm, payload) = multihash::decode_public_key(&bytes)?;
-        PublicKey::from_bytes(algorithm, &payload).map(|key| *key.0)
+        PublicKeyInner::from_bytes(algorithm, &payload)
     }
 }
 
@@ -362,6 +364,22 @@ impl PublicKeyInner {
 }
 
 impl PublicKeyInner {
+    fn from_bytes(algorithm: Algorithm, payload: &[u8]) -> Result<Self, ParseError> {
+        match algorithm {
+            Algorithm::Ed25519 => {
+                ed25519::Ed25519Sha512::parse_public_key(payload).map(PublicKeyInner::Ed25519)
+            }
+            Algorithm::Secp256k1 => secp256k1::EcdsaSecp256k1Sha256::parse_public_key(payload)
+                .map(PublicKeyInner::Secp256k1),
+            Algorithm::BlsNormal => {
+                bls::BlsNormal::parse_public_key(payload).map(PublicKeyInner::BlsNormal)
+            }
+            Algorithm::BlsSmall => {
+                bls::BlsSmall::parse_public_key(payload).map(PublicKeyInner::BlsSmall)
+            }
+        }
+    }
+
     fn to_raw(&self) -> (Algorithm, Vec<u8>) {
         (self.algorithm(), self.payload())
     }
@@ -388,6 +406,109 @@ impl PublicKeyInner {
     }
 }
 
+/// `PublicKey` will be lazily deserialized inside WASM.
+/// This is needed for performance reasons, since `PublicKeyInner::from_bytes` is quite slow.
+/// However inside WASM in most cases `PublicKey` is used only for comparisons (==).
+/// See https://github.com/hyperledger/iroha/issues/5038 for details.
+#[cfg(target_family = "wasm")]
+mod lazy {
+    use alloc::{boxed::Box, vec::Vec};
+    use core::{borrow::Borrow, cell::OnceCell, fmt};
+
+    use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+    use crate::{Algorithm, PublicKeyInner};
+
+    #[derive(Clone, Eq)]
+    pub struct PublicKeyLazy {
+        algorithm: Algorithm,
+        payload: Vec<u8>,
+        inner: OnceCell<Box<PublicKeyInner>>,
+    }
+
+    impl PublicKeyLazy {
+        pub fn new(inner: PublicKeyInner) -> Self {
+            Self {
+                algorithm: inner.algorithm(),
+                payload: inner.payload(),
+                inner: OnceCell::from(Box::new(inner)),
+            }
+        }
+
+        pub fn new_lazy(algorithm: Algorithm, payload: Vec<u8>) -> Self {
+            Self {
+                algorithm,
+                payload,
+                inner: OnceCell::new(),
+            }
+        }
+
+        fn get_inner(&self) -> &PublicKeyInner {
+            self.inner.get_or_init(|| {
+                let inner = PublicKeyInner::from_bytes(self.algorithm, &self.payload)
+                    .expect("Public key deserialization at WASM side must not fail because data received from host side");
+                Box::new(inner)
+            })
+        }
+
+        pub fn algorithm(&self) -> Algorithm {
+            self.algorithm
+        }
+
+        pub fn to_raw(&self) -> (Algorithm, Vec<u8>) {
+            (self.algorithm, self.payload.clone())
+        }
+    }
+
+    impl Borrow<PublicKeyInner> for PublicKeyLazy {
+        fn borrow(&self) -> &PublicKeyInner {
+            self.get_inner()
+        }
+    }
+
+    impl fmt::Display for PublicKeyLazy {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            fmt::Display::fmt(self.get_inner(), f)
+        }
+    }
+
+    impl fmt::Debug for PublicKeyLazy {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            fmt::Debug::fmt(self.get_inner(), f)
+        }
+    }
+
+    impl<'de> Deserialize<'de> for PublicKeyLazy {
+        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+        where
+            D: Deserializer<'de>,
+        {
+            let inner = PublicKeyInner::deserialize(deserializer)?;
+            Ok(PublicKeyLazy::new(inner))
+        }
+    }
+
+    impl Serialize for PublicKeyLazy {
+        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            self.get_inner().serialize(serializer)
+        }
+    }
+
+    impl PartialEq for PublicKeyLazy {
+        fn eq(&self, other: &Self) -> bool {
+            self.algorithm == other.algorithm && self.payload == other.payload
+        }
+    }
+}
+
+#[cfg(not(target_family = "wasm"))]
+type PublicKeyInnerType = Box<PublicKeyInner>;
+#[cfg(target_family = "wasm")]
+type PublicKeyInnerType = PublicKeyLazy;
+
 ffi::ffi_item! {
     /// Public key used in signatures.
     ///
@@ -414,32 +535,32 @@ ffi::ffi_item! {
     #[cfg_attr(not(feature="ffi_import"), display(fmt = "{_0}"))]
     #[cfg_attr(all(feature = "ffi_export", not(feature = "ffi_import")), ffi_type(opaque))]
     #[allow(missing_docs)]
-    pub struct PublicKey(Box<PublicKeyInner>);
+    pub struct PublicKey(PublicKeyInnerType);
 }
 
 #[ffi_impl_opaque]
 impl PublicKey {
+    #[cfg(not(target_family = "wasm"))]
+    fn new(inner: PublicKeyInner) -> Self {
+        Self(Box::new(inner))
+    }
+    #[cfg(target_family = "wasm")]
+    fn new(inner: PublicKeyInner) -> Self {
+        Self(PublicKeyLazy::new(inner))
+    }
+
     /// Creates a new public key from raw bytes received from elsewhere
     ///
     /// # Errors
     ///
     /// Fails if public key parsing fails
     pub fn from_bytes(algorithm: Algorithm, payload: &[u8]) -> Result<Self, ParseError> {
-        match algorithm {
-            Algorithm::Ed25519 => {
-                ed25519::Ed25519Sha512::parse_public_key(payload).map(PublicKeyInner::Ed25519)
-            }
-            Algorithm::Secp256k1 => secp256k1::EcdsaSecp256k1Sha256::parse_public_key(payload)
-                .map(PublicKeyInner::Secp256k1),
-            Algorithm::BlsNormal => {
-                bls::BlsNormal::parse_public_key(payload).map(PublicKeyInner::BlsNormal)
-            }
-            Algorithm::BlsSmall => {
-                bls::BlsSmall::parse_public_key(payload).map(PublicKeyInner::BlsSmall)
-            }
-        }
-        .map(Box::new)
-        .map(PublicKey)
+        #[cfg(not(target_family = "wasm"))]
+        let inner = Box::new(PublicKeyInner::from_bytes(algorithm, payload)?);
+        #[cfg(target_family = "wasm")]
+        let inner = PublicKeyLazy::new_lazy(algorithm, payload.to_vec());
+
+        Ok(Self(inner))
     }
 
     /// Extracts raw bytes from the public key, copying the payload.
@@ -548,7 +669,7 @@ impl FromStr for PublicKey {
     type Err = ParseError;
 
     fn from_str(key: &str) -> Result<Self, Self::Err> {
-        PublicKeyInner::from_str(key).map(Box::new).map(Self)
+        Ok(Self::new(PublicKeyInner::from_str(key)?))
     }
 }
 
@@ -577,7 +698,7 @@ impl From<PrivateKey> for PublicKey {
             (BlsSmall, bls::BlsSmall)
         );
 
-        Self(Box::new(inner))
+        Self::new(inner)
     }
 }