refactor(keymap): create KeyMap type

- creates a new `key_map` module, and `KeyMap` type to replace the
  existing `BTreeMap` alias.
- it's a commit extracted from superseded PR rust-bitcoin#765.

Author: tcharding

src/descriptor/key_map.rs

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: CC0-1.0
+
+//! A map of public key to secret key.
+
+use core::iter;
+
+use bitcoin::secp256k1::{Secp256k1, Signing};
+
+#[cfg(doc)]
+use super::Descriptor;
+use super::{DescriptorKeyParseError, DescriptorPublicKey, DescriptorSecretKey};
+use crate::prelude::{btree_map, BTreeMap};
+
+/// Alias type for a map of public key to secret key.
+///
+/// This map is returned whenever a descriptor that contains secrets is parsed using
+/// [`Descriptor::parse_descriptor`], since the descriptor will always only contain
+/// public keys. This map allows looking up the corresponding secret key given a
+/// public key from the descriptor.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct KeyMap {
+    map: BTreeMap<DescriptorPublicKey, DescriptorSecretKey>,
+}
+
+impl KeyMap {
+    /// Creates a new empty `KeyMap`.
+    #[inline]
+    pub fn new() -> Self { Self { map: BTreeMap::new() } }
+
+    /// Inserts secret key into key map returning the associated public key.
+    #[inline]
+    pub fn insert<C: Signing>(
+        &mut self,
+        secp: &Secp256k1<C>,
+        sk: DescriptorSecretKey,
+    ) -> Result<DescriptorPublicKey, DescriptorKeyParseError> {
+        let pk = sk.to_public(secp)?;
+        if !self.map.contains_key(&pk) {
+            self.map.insert(pk.clone(), sk);
+        }
+        Ok(pk)
+    }
+
+    /// Gets the secret key associated with `pk` if `pk` is in the map.
+    #[inline]
+    pub fn get(&self, pk: &DescriptorPublicKey) -> Option<&DescriptorSecretKey> { self.map.get(pk) }
+
+    /// Returns the number of items in this map.
+    #[inline]
+    pub fn len(&self) -> usize { self.map.len() }
+
+    /// Returns true if the map is empty.
+    #[inline]
+    pub fn is_empty(&self) -> bool { self.map.is_empty() }
+}
+
+impl Default for KeyMap {
+    fn default() -> Self { Self::new() }
+}
+
+impl IntoIterator for KeyMap {
+    type Item = (DescriptorPublicKey, DescriptorSecretKey);
+    type IntoIter = btree_map::IntoIter<DescriptorPublicKey, DescriptorSecretKey>;
+
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter { self.map.into_iter() }
+}
+
+impl iter::Extend<(DescriptorPublicKey, DescriptorSecretKey)> for KeyMap {
+    #[inline]
+    fn extend<T>(&mut self, iter: T)
+    where
+        T: IntoIterator<Item = (DescriptorPublicKey, DescriptorSecretKey)>,
+    {
+        self.map.extend(iter)
+    }
+}

src/descriptor/mod.rs

@@ -51,20 +51,14 @@ pub use self::tr::{
 
 pub mod checksum;
 mod key;
+mod key_map;
 
 pub use self::key::{
     DefiniteDescriptorKey, DerivPaths, DescriptorKeyParseError, DescriptorMultiXKey,
     DescriptorPublicKey, DescriptorSecretKey, DescriptorXKey, InnerXKey, MalformedKeyDataKind,
     NonDefiniteKeyError, SinglePriv, SinglePub, SinglePubKey, Wildcard, XKeyNetwork,
 };
-
-/// Alias type for a map of public key to secret key
-///
-/// This map is returned whenever a descriptor that contains secrets is parsed using
-/// [`Descriptor::parse_descriptor`], since the descriptor will always only contain
-/// public keys. This map allows looking up the corresponding secret key given a
-/// public key from the descriptor.
-pub type KeyMap = BTreeMap<DescriptorPublicKey, DescriptorSecretKey>;
+pub use self::key_map::KeyMap;
 
 /// Script descriptor
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -725,28 +719,24 @@ impl Descriptor<DescriptorPublicKey> {
             key_map: &mut KeyMap,
             secp: &secp256k1::Secp256k1<C>,
         ) -> Result<DescriptorPublicKey, Error> {
-            let (public_key, secret_key) = match DescriptorSecretKey::from_str(s) {
-                Ok(sk) => (
-                    sk.to_public(secp)
-                        .map_err(|e| Error::Unexpected(e.to_string()))?,
-                    Some(sk),
-                ),
-                Err(_) => (
+            match DescriptorSecretKey::from_str(s) {
+                Ok(sk) => {
+                    let pk = key_map
+                        .insert(secp, sk)
+                        .map_err(|e| Error::Unexpected(e.to_string()))?;
+                    Ok(pk)
+                }
+                Err(_) => {
                     // try to parse as a public key if parsing as a secret key failed
-                    s.parse()
-                        .map_err(|e| Error::Parse(ParseError::box_from_str(e)))?,
-                    None,
-                ),
-            };
-
-            if let Some(secret_key) = secret_key {
-                key_map.insert(public_key.clone(), secret_key);
+                    let pk = s
+                        .parse()
+                        .map_err(|e| Error::Parse(ParseError::box_from_str(e)))?;
+                    Ok(pk)
+                }
             }
-
-            Ok(public_key)
         }
 
-        let mut keymap_pk = KeyMapWrapper(BTreeMap::new(), secp);
+        let mut keymap_pk = KeyMapWrapper(KeyMap::new(), secp);
 
         struct KeyMapWrapper<'a, C: secp256k1::Signing>(KeyMap, &'a secp256k1::Secp256k1<C>);
 

src/lib.rs

@@ -777,7 +777,7 @@ mod prelude {
     pub use alloc::{
         borrow::{Borrow, Cow, ToOwned},
         boxed::Box,
-        collections::{vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap},
+        collections::{btree_map, vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap},
         rc, slice,
         string::{String, ToString},
         sync,
@@ -787,7 +787,9 @@ mod prelude {
     pub use std::{
         borrow::{Borrow, Cow, ToOwned},
         boxed::Box,
-        collections::{vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet},
+        collections::{
+            btree_map, vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet,
+        },
         rc, slice,
         string::{String, ToString},
         sync,