Add taproot branch information to the derivation result (#1274)

Author: NicolasDorier

NBitcoin.Tests/MiniscriptTests.cs

@@ -524,7 +524,7 @@ public void CanMassDeriveMiniscripts(AddressIntent intent)
 			var expected = $"multi(2,{a},{b},{c})";
 			Assert.Equal(expected, derived.ToString());
 
-			var derivedKey = allMiniscripts[3].DerivedKeys[keys[0]];
+			var derivedKey = allMiniscripts[3].DerivedKeys[0];
 			Assert.Equal(new KeyPath($"{typeIdx[0]}/{15 + 3}"), derivedKey.KeyPath);
 			Assert.Equal(a, ((Value.PubKeyValue)derivedKey.Pubkey).PubKey);
 		}

NBitcoin/WalletPolicies/DerivationResult.cs

@@ -11,7 +11,7 @@ namespace NBitcoin.WalletPolicies;
 
 public class DerivationResult
 {
-	internal DerivationResult(Miniscript miniscript, Dictionary<HDKeyNode, Derivation> derivations)
+	internal DerivationResult(Miniscript miniscript, List<Derivation> derivations)
 	{
 		Miniscript = miniscript;
 		DerivedKeys = derivations;
@@ -21,20 +21,19 @@ internal DerivationResult(Miniscript miniscript, Dictionary<HDKeyNode, Derivatio
 	/// <summary>
 	/// The derived keys. The values are either <see cref="Value.PubKeyValue"/> or a <see cref="Value.TaprootPubKeyValue"/>.
 	/// </summary>
-	public Dictionary<HDKeyNode, Derivation> DerivedKeys { get; }
+	public List<Derivation> DerivedKeys { get; }
 }
 
-public class Derivation
-{
-	public Derivation(KeyPath keyPath, Value pubkey)
-	{
-		KeyPath = keyPath;
-		Pubkey = pubkey;
-	}
-	public KeyPath KeyPath { get; }
-	/// <summary>
-	/// The derived key. This is either a <see cref="Value.PubKeyValue"/> or a <see cref="Value.TaprootPubKeyValue"/>.
-	/// </summary>
-	public Value Pubkey { get; }
-}
+/// <summary>
+/// Represents a derivation that has been made from a <see cref="HDKeyNode"/>.
+/// </summary>
+/// <param name="KeyPath">The derived key path.</param>
+/// <param name="Pubkey">The derived public key (either <see cref="Value.PubKeyValue"/> or <see cref="Value.TaprootPubKeyValue"/>).</param>
+/// <param name="TaprootBranch">The taproot branch in which this has been derived, if any.</param>
+/// <param name="Source">The source node that generated this key.</param>
+public record Derivation(
+	KeyPath KeyPath,
+	Value Pubkey,
+	TaprootBranchNode? TaprootBranch,
+	HDKeyNode Source);
 #endif

NBitcoin/WalletPolicies/Visitors/DeriveVisitor.cs

@@ -32,8 +32,22 @@ public DerivationResult[] Derive(MiniscriptNode node, Network network)
 	}
 
 	internal static readonly byte[] BIP0328CC = Encoders.Hex.DecodeData("868087ca02a6f974c4598924c36b57762d32cb45717167e300622c7167e38965");
+	Stack<TaprootBranchNode> _TaprootBranches = new();
+	TaprootBranchNode? TaprootBranch => _TaprootBranches.Count > 0 ? _TaprootBranches.Peek() : null;
 	public override MiniscriptNode Visit(MiniscriptNode node)
 	{
+		if (node is TaprootBranchNode b)
+		{
+			_TaprootBranches.Push(b);
+			try
+			{
+				return base.Visit(b);
+			}
+			finally
+			{
+				_TaprootBranches.Pop();
+			}
+		}
 		if (node is MultipathNode { Target: MusigNode })
 		{
 			var wasNestedMusig = _nestedMusig;
@@ -53,10 +67,10 @@ public override MiniscriptNode Visit(MiniscriptNode node)
 		}
 		if (node is MiniscriptNode.MultipathNode mki && mki.CanDerive(Intent))
 		{
-			if (mki.Target is HDKeyNode { Key: var pk } xpub)
+			if (mki.Target is HDKeyNode xpub)
 			{
-				var value = GetPublicKey(mki, pk);
-				_Derivations.TryAdd(xpub, value);
+				var value = GetPublicKey(mki, xpub.Key, xpub);
+				_Derivations.Add(new(value.KeyPath, value.Pubkey, TaprootBranch, xpub));
 				node = value.Pubkey;
 			}
 			else if (mki.Target is MusigNode musig)
@@ -69,19 +83,21 @@ public override MiniscriptNode Visit(MiniscriptNode node)
 		return node;
 	}
 
-	private Derivation GetPublicKey(MiniscriptNode.MultipathNode mki, IHDKey k)
+	private (KeyPath KeyPath, Value Pubkey) GetPublicKey(MiniscriptNode.MultipathNode mki, IHDKey k, HDKeyNode? source = null)
 	{
 		var type = mki.GetTypeIndex(Intent);
 		k = DeriveIntent(k, type);
 		k = k.Derive((uint)idx);
 		var keyType = _nestedMusig ? KeyType.Classic : KeyType;
-		return new Derivation(new KeyPath([(uint)type, (uint)idx]), keyType switch
-		{
-			KeyType.Taproot => MiniscriptNode.Create(k.GetPublicKey().TaprootPubKey),
-			_ => MiniscriptNode.Create(k.GetPublicKey())
-		});
+		return (
+			new KeyPath([(uint)type, (uint)idx]),
+			keyType switch
+			{
+				KeyType.Taproot => MiniscriptNode.Create(k.GetPublicKey().TaprootPubKey),
+				_ => MiniscriptNode.Create(k.GetPublicKey())
+			});
 	}
-	Dictionary<HDKeyNode, Derivation> _Derivations = new();
+	List<Derivation> _Derivations = new();
 
 	public bool _nestedMusig = false;