Balance for p2c

crates/chain/src/keychain.rs

@@ -10,8 +10,9 @@
 //!
 //! [`SpkTxOutIndex`]: crate::SpkTxOutIndex
 
+/// Indexer for TxOut
 #[cfg(feature = "miniscript")]
-mod txout_index;
+pub mod txout_index;
 use tapyrus::Amount;
 #[cfg(feature = "miniscript")]
 pub use txout_index::*;

crates/chain/src/keychain/txout_index.rs

@@ -10,8 +10,8 @@ use core::{
     ops::{Bound, RangeBounds},
 };
 use tapyrus::{
-    hashes::Hash, script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, Script,
-    SignedAmount, Transaction, TxOut,
+    hashes::Hash, script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint,
+    PublicKey, Script, ScriptBuf, SignedAmount, Transaction, TxOut,
 };
 
 use crate::Append;
@@ -378,6 +378,11 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         Some((keychain.clone(), *last_index))
     }
 
+    /// Insert payment base key for pay-to-contract script pubkey
+    pub fn insert_p2c_spk(&mut self, spk: ScriptBuf, payment_base: PublicKey) {
+        self.inner.insert_p2c_spk(spk, payment_base);
+    }
+
     /// Returns whether the spk under the `keychain`'s `index` has been used.
     ///
     /// Here, "unused" means that after the script pubkey was stored in the index, the index has

crates/chain/src/spk_txout_index.rs

@@ -5,8 +5,8 @@ use crate::{
     indexed_tx_graph::Indexer,
 };
 use tapyrus::{
-    script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, Script, ScriptBuf,
-    SignedAmount, Transaction, TxOut,
+    script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, PublicKey, Script,
+    ScriptBuf, SignedAmount, Transaction, TxOut,
 };
 
 /// An index storing [`TxOut`]s that have a script pubkey that matches those in a list.
@@ -41,6 +41,8 @@ pub struct SpkTxOutIndex<I> {
     txouts: BTreeMap<OutPoint, (I, TxOut)>,
     /// Lookup from spk index to outpoints that had that spk
     spk_txouts: BTreeSet<(I, OutPoint)>,
+    /// Pay-to-contract payment_base lookup by p2c spk
+    p2c_spks: HashMap<ScriptBuf, ScriptBuf>,
 }
 
 impl<I> Default for SpkTxOutIndex<I> {
@@ -51,6 +53,7 @@ impl<I> Default for SpkTxOutIndex<I> {
             spk_indices: Default::default(),
             spk_txouts: Default::default(),
             unused: Default::default(),
+            p2c_spks: Default::default(),
         }
     }
 }
@@ -103,12 +106,17 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
     /// Scan a single `TxOut` for a matching script pubkey and returns the index that matches the
     /// script pubkey (if any).
     pub fn scan_txout(&mut self, op: OutPoint, txout: &TxOut) -> Option<&I> {
-        let spk_i = if txout.script_pubkey.is_colored() {
-            self.spk_indices.get(&ScriptBuf::from_bytes(
-                txout.script_pubkey.as_bytes()[35..].to_vec(),
-            ))
+        let script_pubkey = if txout.script_pubkey.is_colored() {
+            txout.script_pubkey.remove_color()
         } else {
-            self.spk_indices.get(&txout.script_pubkey)
+            txout.script_pubkey.clone()
+        };
+        let payment_base = self.p2c_spks.get(&script_pubkey);
+
+        let spk_i = if let Some(p) = payment_base {
+            self.spk_indices.get(p.as_script())
+        } else {
+            self.spk_indices.get(&script_pubkey)
         };
         if let Some(spk_i) = spk_i {
             self.txouts.insert(op, (spk_i.clone(), txout.clone()));
@@ -192,6 +200,12 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
         &self.spks
     }
 
+    /// Insert payment base key for pay-to-contract script pubkey
+    pub fn insert_p2c_spk(&mut self, spk: ScriptBuf, payment_base: PublicKey) {
+        let p2c_spk = ScriptBuf::new_p2pkh(&payment_base.pubkey_hash());
+        self.p2c_spks.insert(spk, p2c_spk);
+    }
+
     /// Adds a script pubkey to scan for. Returns `false` and does nothing if spk already exists in the map
     ///
     /// the index will look for outputs spending to this spk whenever it scans new data.
@@ -304,11 +318,17 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
         }
         for txout in &tx.output {
             let script_pubkey = if txout.script_pubkey.is_colored() {
-                ScriptBuf::from_bytes(txout.script_pubkey.as_bytes()[35..].to_vec())
+                txout.script_pubkey.remove_color()
             } else {
                 txout.script_pubkey.clone()
             };
-            if let Some(index) = self.index_of_spk(&script_pubkey) {
+            let payment_base = self.p2c_spks.get(&script_pubkey);
+            let script_pubkey_ref = if let Some(p) = payment_base {
+                p
+            } else {
+                &script_pubkey
+            };
+            if let Some(index) = self.index_of_spk(script_pubkey_ref) {
                 if range.contains(index)
                     && txout.script_pubkey.color_id().unwrap_or_default() == *color_id
                 {

crates/sqlite/src/store.rs

@@ -502,7 +502,7 @@ impl<K, A> Store<K, A> {
             let payment_base: Vec<u8> = c.payment_base.to_bytes();
             let spendable: u32 = if c.spendable { 1 } else { 0 };
             insert_contract_stmt.execute(named_params! {
-                ":contract_id": contract_id, ":contract": contract, ":payment_base": payment_base, ":spendable": spendable})
+                ":contract_id": contract_id, ":contract": contract, ":payment_base": payment_base, ":spendable": spendable })
                 .map_err(Error::Sqlite)?;
         }
         Ok(())
@@ -663,17 +663,17 @@ mod test {
             agg_changeset.unwrap().contract.get("id").unwrap().spendable,
             true
         );
-
+        let payment_base = PublicKey::from_str(
+            "028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
+        )
+        .unwrap();
         let mut contract: contract::ChangeSet = contract::ChangeSet::new();
         contract.insert(
             "id".to_string(),
             Contract {
                 contract_id: "id".to_string(),
                 contract: vec![0x00, 0x01, 0x02],
-                payment_base: PublicKey::from_str(
-                    "028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
-                )
-                .unwrap(),
+                payment_base,
                 spendable: false,
             },
         );
@@ -881,6 +881,11 @@ mod test {
                 indexer: keychain::ChangeSet::default(),
             };
 
+        let payment_base = PublicKey::from_str(
+            "028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
+        )
+        .unwrap();
+
         let mut contract: contract::ChangeSet = contract::ChangeSet::new();
         contract.insert(
             "id".to_string(),
@@ -890,10 +895,7 @@ mod test {
                     0x00, 0x00, 0x55, 0x0e, 0x84, 0x00, 0xe2, 0x9b, 0x41, 0xd4, 0xa7, 0x16, 0x44,
                     0x66, 0x55, 0x44, 0x00, 0x00,
                 ],
-                payment_base: PublicKey::from_str(
-                    "028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
-                )
-                .unwrap(),
+                payment_base,
                 spendable: true,
             },
         );