From e7f6a601f91dbb7f1c2dc5d5abf5d1c176a557a9 Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Sat, 9 May 2026 16:21:34 -0600
Subject: [PATCH] ml-dsa: precompute `VerifyingKey`s when `alloc` enabled
 (#1346)

Adds an `alloc`-gated `verifying_key` field to `SigningKey`, and when
the feature is enabled precomputes the key at the time the `SigningKey`
is initialized.

We previously used to do this but stopped to optimize stack usage (see
the changes in #1259 and #1261), however when `alloc` is enabled this
isn't an issue since we've moved the relevant data to the heap in this
case (see #1344 and #1345), so stack usage is no longer an issue.

This makes it possible to implement `signature::KeypairRef` for
`SigningKey` which still provides a blanket `signature::Keypair` impl,
so we can always depend on the latter being availble but take advantage
of the former when `alloc` is enabled.
---
 ml-dsa/src/signing.rs | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/ml-dsa/src/signing.rs b/ml-dsa/src/signing.rs
index d9728d7..22cd562 100644
--- a/ml-dsa/src/signing.rs
+++ b/ml-dsa/src/signing.rs
@@ -41,6 +41,10 @@ pub struct SigningKey<P: MlDsaParams> {
 
     /// The seed this signing key was derived from
     seed: MaybeBox<Seed>,
+
+    /// When the `alloc` feature is available, precompute the [`VerifyingKey`].
+    #[cfg(feature = "alloc")]
+    verifying_key: VerifyingKey<P>,
 }
 
 impl<P: MlDsaParams> SigningKey<P> {
@@ -73,11 +77,16 @@ impl<P: MlDsaParams> SigningKey<P> {
 
         let enc = VerifyingKey::<P>::encode_internal(&rho, &t1);
         let tr: B64 = H::default().absorb(&enc).squeeze_new();
-        let signing_key = ExpandedSigningKey::new(rho, K, tr, s1, s2, t0, A_hat);
+        let expanded_key = ExpandedSigningKey::new(rho, K, tr, s1, s2, t0, A_hat);
+
+        #[cfg(feature = "alloc")]
+        let verifying_key = expanded_key.verifying_key();
 
         SigningKey {
-            expanded_key: MaybeBox::new(signing_key),
+            expanded_key: MaybeBox::new(expanded_key),
             seed: MaybeBox::new(xi.clone()),
+            #[cfg(feature = "alloc")]
+            verifying_key,
         }
     }
 
@@ -134,6 +143,10 @@ impl<P: MlDsaParams> fmt::Debug for SigningKey<P> {
     }
 }
 
+// NOTE: when the `alloc` feature is enabled, we receive a blanket impl of `Keypair` via the impl
+// of the `KeypairRef` trait which simply clones the precomputed verifying key, providing equivalent
+// functionality and thus this is actually still an additive use of features.
+#[cfg(not(feature = "alloc"))]
 impl<P: MlDsaParams> signature::Keypair for SigningKey<P> {
     type VerifyingKey = VerifyingKey<P>;
     fn verifying_key(&self) -> VerifyingKey<P> {
@@ -141,6 +154,18 @@ impl<P: MlDsaParams> signature::Keypair for SigningKey<P> {
     }
 }
 
+#[cfg(feature = "alloc")]
+impl<P: MlDsaParams> AsRef<VerifyingKey<P>> for SigningKey<P> {
+    fn as_ref(&self) -> &VerifyingKey<P> {
+        &self.verifying_key
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<P: MlDsaParams> signature::KeypairRef for SigningKey<P> {
+    type VerifyingKey = VerifyingKey<P>;
+}
+
 /// The `Signer` implementation for `SigningKey` uses the optional deterministic variant of ML-DSA, and
 /// only supports signing with an empty context string.
 impl<P: MlDsaParams> Signer<Signature<P>> for SigningKey<P> {