Add KeyExpr inside Taproot descriptor

kanishk779 · kanishk779 · commit 297c954dc16a · 2022-08-24T10:05:15.000+05:30
diff --git a/examples/taproot.rs b/examples/taproot.rs
@@ -75,7 +75,7 @@ fn main() {
     if let Descriptor::Tr(ref p) = desc {
         // Check if internal key is correctly inferred as Ca
         // assert_eq!(p.internal_key(), &pubkeys[2]);
-        assert_eq!(p.internal_key(), "Ca");
+        assert_eq!(p.internal_key().single_key().unwrap(), "Ca");
 
         // Iterate through scripts
         let mut iter = p.iter_scripts();
diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs
@@ -271,7 +271,7 @@ impl<Pk: MiniscriptKey> Descriptor<Pk> {
 
     /// Create new tr descriptor
     /// Errors when miniscript exceeds resource limits under Tap context
-    pub fn new_tr(key: Pk, script: Option<tr::TapTree<Pk>>) -> Result<Self, Error> {
+    pub fn new_tr(key: KeyExpr<Pk>, script: Option<tr::TapTree<Pk>>) -> Result<Self, Error> {
         Ok(Descriptor::Tr(Tr::new(key, script)?))
     }
 
diff --git a/src/descriptor/tr.rs b/src/descriptor/tr.rs
@@ -13,6 +13,7 @@ use sync::Arc;
 
 use super::checksum::{desc_checksum, verify_checksum};
 use crate::expression::{self, FromTree};
+use crate::miniscript::musig_key::KeyExpr;
 use crate::miniscript::Miniscript;
 use crate::policy::semantic::Policy;
 use crate::policy::Liftable;
@@ -39,7 +40,7 @@ pub enum TapTree<Pk: MiniscriptKey> {
 /// A taproot descriptor
 pub struct Tr<Pk: MiniscriptKey> {
     /// A taproot internal key
-    internal_key: Pk,
+    internal_key: KeyExpr<Pk>,
     /// Optional Taproot Tree with spending conditions
     tree: Option<TapTree<Pk>>,
     /// Optional spending information associated with the descriptor
@@ -163,9 +164,8 @@ impl<Pk: MiniscriptKey> fmt::Debug for TapTree<Pk> {
 
 impl<Pk: MiniscriptKey> Tr<Pk> {
     /// Create a new [`Tr`] descriptor from internal key and [`TapTree`]
-    pub fn new(internal_key: Pk, tree: Option<TapTree<Pk>>) -> Result<Self, Error> {
+    pub fn new(internal_key: KeyExpr<Pk>, tree: Option<TapTree<Pk>>) -> Result<Self, Error> {
         let nodes = tree.as_ref().map(|t| t.taptree_height()).unwrap_or(0);
-
         if nodes <= TAPROOT_CONTROL_MAX_NODE_COUNT {
             Ok(Self {
                 internal_key,
@@ -186,7 +186,7 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
     }
 
     /// Obtain the internal key of [`Tr`] descriptor
-    pub fn internal_key(&self) -> &Pk {
+    pub fn internal_key(&self) -> &KeyExpr<Pk> {
         &self.internal_key
     }
 
@@ -226,7 +226,7 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
         let secp = secp256k1::Secp256k1::verification_only();
         // Key spend path with no merkle root
         let data = if self.tree.is_none() {
-            TaprootSpendInfo::new_key_spend(&secp, self.internal_key.to_x_only_pubkey(), None)
+            TaprootSpendInfo::new_key_spend(&secp, self.internal_key.key_agg(), None)
         } else {
             let mut builder = TaprootBuilder::new();
             for (depth, ms) in self.iter_scripts() {
@@ -236,7 +236,7 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
                     .expect("Computing spend data on a valid Tree should always succeed");
             }
             // Assert builder cannot error here because we have a well formed descriptor
-            match builder.finalize(&secp, self.internal_key.to_x_only_pubkey()) {
+            match builder.finalize(&secp, self.internal_key.key_agg()) {
                 Ok(data) => data,
                 Err(e) => match e {
                     TaprootBuilderError::InvalidMerkleTreeDepth(_) => {
@@ -419,7 +419,7 @@ impl_from_tree!(
                             key.args.len()
                         )));
                     }
-                    Tr::new(expression::terminal(key, Pk::from_str)?, None)
+                    Tr::new(expression::terminal(key, KeyExpr::<Pk>::from_str)?, None)
                 }
                 2 => {
                     let key = &top.args[0];
@@ -431,7 +431,10 @@ impl_from_tree!(
                     }
                     let tree = &top.args[1];
                     let ret = Self::parse_tr_script_spend(tree)?;
-                    Tr::new(expression::terminal(key, Pk::from_str)?, Some(ret))
+                    Tr::new(
+                        expression::terminal(key, KeyExpr::<Pk>::from_str)?,
+                        Some(ret),
+                    )
                 }
                 _ => {
                     return Err(Error::Unexpected(format!(
@@ -567,12 +570,10 @@ impl<Pk: MiniscriptKey> Liftable<Pk> for Tr<Pk> {
         match &self.tree {
             Some(root) => Ok(Policy::Threshold(
                 1,
-                vec![
-                    Policy::KeyHash(self.internal_key.to_pubkeyhash()),
-                    root.lift()?,
-                ],
+                vec![self.internal_key.lift()?, root.lift()?],
             )),
-            None => Ok(Policy::KeyHash(self.internal_key.to_pubkeyhash())),
+            // None => Ok(Policy::KeyHash(self.internal_key.to_pubkeyhash())),
+            None => self.internal_key.lift(),
         }
     }
 }
@@ -586,7 +587,7 @@ impl<Pk: MiniscriptKey> ForEachKey<Pk> for Tr<Pk> {
         let script_keys_res = self
             .iter_scripts()
             .all(|(_d, ms)| ms.for_each_key(&mut pred));
-        script_keys_res && pred(&self.internal_key)
+        script_keys_res && self.internal_key().for_any_key(pred)
     }
 }
 
@@ -602,7 +603,7 @@ where
         T: Translator<P, Q, E>,
     {
         let translate_desc = Tr {
-            internal_key: translate.pk(&self.internal_key)?,
+            internal_key: self.internal_key.translate_pk(translate)?,
             tree: match &self.tree {
                 Some(tree) => Some(tree.translate_helper(translate)?),
                 None => None,
diff --git a/src/miniscript/mod.rs b/src/miniscript/mod.rs
@@ -519,11 +519,10 @@ mod tests {
         let desc = Descriptor::<String>::from_str("tr(X,{pk(musig(X1)),multi_a(1,X2,X3)})");
         assert_eq!(desc.is_ok(), true);
 
-        let desc =
-            Descriptor::<String>::from_str("tr(pk(musig(E)),{pk(A),multi_a(1,B,musig(C,D))})");
+        let desc = Descriptor::<String>::from_str("tr(musig(E),{pk(A),multi_a(1,B,musig(C,D))})");
         assert_eq!(desc.is_ok(), true);
 
-        let desc = Descriptor::<String>::from_str("tr(pk(musig(D)),pk(musig(A,B,musig(C))))");
+        let desc = Descriptor::<String>::from_str("tr(musig(D),pk(musig(A,B,musig(C))))");
         assert_eq!(desc.is_ok(), true);
     }
 
diff --git a/src/miniscript/musig_key.rs b/src/miniscript/musig_key.rs
@@ -7,6 +7,8 @@ use bitcoin::secp256k1::{Secp256k1, VerifyOnly};
 use secp256k1_zkp::MusigKeyAggCache;
 
 use crate::expression::{FromTree, Tree};
+use crate::policy::semantic::{Policy, Policy as Semantic};
+use crate::policy::Liftable;
 use crate::prelude::*;
 use crate::{expression, Error, ForEachKey, MiniscriptKey, ToPublicKey, TranslatePk, Translator};
 
@@ -50,6 +52,23 @@ where
     }
 }
 
+impl<Pk: MiniscriptKey> Liftable<Pk> for KeyExpr<Pk> {
+    fn lift(&self) -> Result<Policy<Pk>, Error> {
+        let res = match self {
+            KeyExpr::SingleKey(pk) => Semantic::KeyHash(pk.to_pubkeyhash()),
+            KeyExpr::MuSig(keys) => {
+                let mut policy_vec: Vec<Semantic<Pk>> = vec![];
+                for key in keys {
+                    policy_vec.push((key.clone()).lift()?)
+                }
+                Semantic::Threshold(keys.len(), policy_vec)
+            }
+        }
+        .normalized();
+        Ok(res)
+    }
+}
+
 #[derive(Debug, Clone)]
 /// Iterator for [`KeyExpr`]
 pub struct KeyExprIter<'a, Pk: MiniscriptKey> {
diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs
@@ -22,6 +22,7 @@ use std::error;
 #[cfg(feature = "compiler")]
 use {
     crate::descriptor::TapTree,
+    crate::miniscript::musig_key::KeyExpr,
     crate::miniscript::ScriptContext,
     crate::policy::compiler::CompilerError,
     crate::policy::compiler::OrdF64,
@@ -277,7 +278,8 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
             _ => {
                 let (internal_key, policy) = self.clone().extract_key(unspendable_key)?;
                 let tree = Descriptor::new_tr(
-                    internal_key,
+                    // Temporary solution, need to decide what we should write in place of singlekey
+                    KeyExpr::SingleKey(internal_key),
                     match policy {
                         Policy::Trivial => None,
                         policy => {
diff --git a/src/policy/mod.rs b/src/policy/mod.rs
@@ -256,6 +256,8 @@ mod tests {
 
     use super::super::miniscript::context::{Segwitv0, Tap};
     use super::super::miniscript::Miniscript;
+    #[cfg(feature = "compiler")]
+    use super::KeyExpr;
     use super::{Concrete, Liftable, Semantic};
     use crate::prelude::*;
     use crate::DummyKey;
@@ -448,7 +450,8 @@ mod tests {
             let ms_compilation: Miniscript<String, Tap> = ms_str!("multi_a(2,A,B,C,D)");
             let tree: TapTree<String> = TapTree::Leaf(Arc::new(ms_compilation));
             let expected_descriptor =
-                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+                Descriptor::new_tr(KeyExpr::SingleKey(unspendable_key.clone()), Some(tree))
+                    .unwrap();
             assert_eq!(descriptor, expected_descriptor);
         }
 
@@ -465,7 +468,8 @@ mod tests {
             let right_node: Arc<TapTree<String>> = Arc::from(TapTree::Leaf(right_ms_compilation));
             let tree: TapTree<String> = TapTree::Tree(left_node, right_node);
             let expected_descriptor =
-                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+                Descriptor::new_tr(KeyExpr::SingleKey(unspendable_key.clone()), Some(tree))
+                    .unwrap();
             assert_eq!(descriptor, expected_descriptor);
         }
 
@@ -548,7 +552,8 @@ mod tests {
                 )),
             );
 
-            let expected_descriptor = Descriptor::new_tr("E".to_string(), Some(tree)).unwrap();
+            let expected_descriptor =
+                Descriptor::new_tr(KeyExpr::SingleKey("E".to_string()), Some(tree)).unwrap();
             assert_eq!(descriptor, expected_descriptor);
         }
     }
diff --git a/src/psbt/mod.rs b/src/psbt/mod.rs
@@ -33,6 +33,7 @@ use bitcoin::{self, EcdsaSighashType, SchnorrSighashType, Script};
 
 use crate::miniscript::iter::PkPkh;
 use crate::miniscript::limits::SEQUENCE_LOCKTIME_DISABLE_FLAG;
+use crate::miniscript::musig_key::KeyExpr;
 use crate::miniscript::satisfy::{After, Older};
 use crate::prelude::*;
 use crate::{
@@ -1017,13 +1018,15 @@ fn update_input_with_descriptor_helper(
             let ik_xpk = tr_xpk.internal_key();
             input.tap_internal_key = Some(ik_derived);
             input.tap_merkle_root = spend_info.merkle_root();
-            input.tap_key_origins.insert(
-                ik_derived,
-                (
-                    vec![],
-                    (ik_xpk.master_fingerprint(), ik_xpk.full_derivation_path()),
-                ),
-            );
+            match ik_xpk {
+                KeyExpr::SingleKey(pk) => {
+                    input.tap_key_origins.insert(
+                        ik_derived,
+                        (vec![], (pk.master_fingerprint(), pk.full_derivation_path())),
+                    );
+                }
+                KeyExpr::MuSig(_) => (), // Need to fix this, when we add full support to KeyExpr
+            };
 
             for ((_depth_der, ms_derived), (_depth, ms)) in
                 tr_derived.iter_scripts().zip(tr_xpk.iter_scripts())
diff --git a/tests/test_desc.rs b/tests/test_desc.rs
@@ -156,7 +156,7 @@ pub fn test_desc_satisfy(
 
             let internal_key_present = x_only_pks
                 .iter()
-                .position(|&x| x.to_public_key() == *tr.internal_key());
+                .position(|&x| x.to_public_key() == *tr.internal_key().single_key().unwrap());
             let internal_keypair = internal_key_present.map(|idx| xonly_keypairs[idx].clone());
             let prevouts = [witness_utxo];
             let prevouts = sighash::Prevouts::All(&prevouts);

Original file line number	Diff line number	Diff line change
`@@ -271,7 +271,7 @@ impl<Pk: MiniscriptKey> Descriptor<Pk> {`
`271`	`271`
`272`	`272`	`/// Create new tr descriptor`
`273`	`273`	`/// Errors when miniscript exceeds resource limits under Tap context`
`274`		`- pub fn new_tr(key: Pk, script: Option<tr::TapTree<Pk>>) -> Result<Self, Error> {`
	`274`	`+ pub fn new_tr(key: KeyExpr<Pk>, script: Option<tr::TapTree<Pk>>) -> Result<Self, Error> {`
`275`	`275`	`Ok(Descriptor::Tr(Tr::new(key, script)?))`
`276`	`276`	`}`
`277`	`277`