Use ScriptBuf over Address where possible

substrate/client/src/networks/bitcoin.rs

@@ -7,19 +7,23 @@ use bitcoin::{
   PubkeyHash, ScriptHash,
   network::Network,
   WitnessVersion, WitnessProgram, ScriptBuf,
-  address::{AddressType, NetworkChecked, Address as BAddressGeneric},
+  address::{AddressType, NetworkChecked, Address as BAddress},
 };
 
-type BAddress = BAddressGeneric<NetworkChecked>;
-
 #[derive(Clone, Eq, Debug)]
-pub struct Address(BAddress);
+pub struct Address(ScriptBuf);
 
 impl PartialEq for Address {
   fn eq(&self, other: &Self) -> bool {
     // Since Serai defines the Bitcoin-address specification as a variant of the script alone,
     // define equivalency as the script alone
-    self.0.script_pubkey() == other.0.script_pubkey()
+    self.0 == other.0
+  }
+}
+
+impl From<Address> for ScriptBuf {
+  fn from(addr: Address) -> ScriptBuf {
+    addr.0
   }
 }
 
@@ -27,10 +31,11 @@ impl FromStr for Address {
   type Err = ();
   fn from_str(str: &str) -> Result<Address, ()> {
     Address::new(
-      BAddressGeneric::from_str(str)
+      BAddress::from_str(str)
         .map_err(|_| ())?
         .require_network(Network::Bitcoin)
-        .map_err(|_| ())?,
+        .map_err(|_| ())?
+        .script_pubkey(),
     )
     .ok_or(())
   }
@@ -38,7 +43,9 @@ impl FromStr for Address {
 
 impl fmt::Display for Address {
   fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-    self.0.fmt(f)
+    BAddress::<NetworkChecked>::from_script(&self.0, Network::Bitcoin)
+      .map_err(|_| fmt::Error)?
+      .fmt(f)
   }
 }
 
@@ -57,45 +64,40 @@ impl TryFrom<Vec<u8>> for Address {
   fn try_from(data: Vec<u8>) -> Result<Address, ()> {
     Ok(Address(match EncodedAddress::decode(&mut data.as_ref()).map_err(|_| ())? {
       EncodedAddress::P2PKH(hash) => {
-        BAddress::p2pkh(PubkeyHash::from_raw_hash(Hash::from_byte_array(hash)), Network::Bitcoin)
+        ScriptBuf::new_p2pkh(&PubkeyHash::from_raw_hash(Hash::from_byte_array(hash)))
       }
       EncodedAddress::P2SH(hash) => {
-        let script_hash = ScriptHash::from_raw_hash(Hash::from_byte_array(hash));
-        let res =
-          BAddress::from_script(&ScriptBuf::new_p2sh(&script_hash), Network::Bitcoin).unwrap();
-        debug_assert_eq!(res.script_hash(), Some(script_hash));
-        res
+        ScriptBuf::new_p2sh(&ScriptHash::from_raw_hash(Hash::from_byte_array(hash)))
+      }
+      EncodedAddress::P2WPKH(hash) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V0, &hash).unwrap())
+      }
+      EncodedAddress::P2WSH(hash) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V0, &hash).unwrap())
+      }
+      EncodedAddress::P2TR(key) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V1, &key).unwrap())
       }
-      EncodedAddress::P2WPKH(hash) => BAddress::from_witness_program(
-        WitnessProgram::new(WitnessVersion::V0, &hash).unwrap(),
-        Network::Bitcoin,
-      ),
-      EncodedAddress::P2WSH(hash) => BAddress::from_witness_program(
-        WitnessProgram::new(WitnessVersion::V0, &hash).unwrap(),
-        Network::Bitcoin,
-      ),
-      EncodedAddress::P2TR(key) => BAddress::from_witness_program(
-        WitnessProgram::new(WitnessVersion::V1, &key).unwrap(),
-        Network::Bitcoin,
-      ),
     }))
   }
 }
 
 fn try_to_vec(addr: &Address) -> Result<Vec<u8>, ()> {
   let witness_program = |addr: &Address| {
-    let script = addr.0.script_pubkey();
-    let program_push = script.as_script().instructions().last().ok_or(())?.map_err(|_| ())?;
+    let program_push = addr.0.as_script().instructions().last().ok_or(())?.map_err(|_| ())?;
     let program = program_push.push_bytes().ok_or(())?.as_bytes();
     Ok::<_, ()>(program.to_vec())
   };
+
+  let parsed_addr =
+    BAddress::<NetworkChecked>::from_script(&addr.0, Network::Bitcoin).map_err(|_| ())?;
   Ok(
-    (match addr.0.address_type() {
+    (match parsed_addr.address_type() {
       Some(AddressType::P2pkh) => {
-        EncodedAddress::P2PKH(*addr.0.pubkey_hash().unwrap().as_raw_hash().as_byte_array())
+        EncodedAddress::P2PKH(*parsed_addr.pubkey_hash().unwrap().as_raw_hash().as_byte_array())
       }
       Some(AddressType::P2sh) => {
-        EncodedAddress::P2SH(*addr.0.script_hash().unwrap().as_raw_hash().as_byte_array())
+        EncodedAddress::P2SH(*parsed_addr.script_hash().unwrap().as_raw_hash().as_byte_array())
       }
       Some(AddressType::P2wpkh) => {
         let program = witness_program(addr)?;
@@ -127,20 +129,8 @@ impl From<Address> for Vec<u8> {
   }
 }
 
-impl From<Address> for BAddress {
-  fn from(addr: Address) -> BAddress {
-    addr.0
-  }
-}
-
-impl AsRef<BAddress> for Address {
-  fn as_ref(&self) -> &BAddress {
-    &self.0
-  }
-}
-
 impl Address {
-  pub fn new(address: BAddress) -> Option<Self> {
+  pub fn new(address: ScriptBuf) -> Option<Self> {
     let res = Self(address);
     if try_to_vec(&res).is_ok() {
       return Some(res);