From 536c500e3ff8fb076951c7a3ab5eed7eca98f837 Mon Sep 17 00:00:00 2001
From: thunderbiscuit <thunderbiscuit@protonmail.com>
Date: Wed, 6 Dec 2023 21:03:53 -0500
Subject: [PATCH] feat: add drain_to method on txbuilder

---
 bdk-ffi/src/bdk.udl   |  2 ++
 bdk-ffi/src/wallet.rs | 38 +++++++++++++++++++-------------------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/bdk-ffi/src/bdk.udl b/bdk-ffi/src/bdk.udl
index 52a3b52f..474c5c7a 100644
--- a/bdk-ffi/src/bdk.udl
+++ b/bdk-ffi/src/bdk.udl
@@ -138,6 +138,8 @@ interface TxBuilder {
 
   TxBuilder drain_wallet();
 
+  TxBuilder drain_to(Script script);
+
   [Throws=BdkError]
   PartiallySignedTransaction finish([ByRef] Wallet wallet);
 };
diff --git a/bdk-ffi/src/wallet.rs b/bdk-ffi/src/wallet.rs
index 61961afc..9bb11829 100644
--- a/bdk-ffi/src/wallet.rs
+++ b/bdk-ffi/src/wallet.rs
@@ -291,7 +291,7 @@ pub struct TxBuilder {
     pub(crate) fee_rate: Option<f32>,
     pub(crate) fee_absolute: Option<u64>,
     pub(crate) drain_wallet: bool,
-    // pub(crate) drain_to: Option<BdkScript>,
+    pub(crate) drain_to: Option<BdkScriptBuf>,
     // pub(crate) rbf: Option<RbfValue>,
     // pub(crate) data: Vec<u8>,
 }
@@ -307,7 +307,7 @@ impl TxBuilder {
             fee_rate: None,
             fee_absolute: None,
             drain_wallet: false,
-            // drain_to: None,
+            drain_to: None,
             // rbf: None,
             // data: Vec::new(),
         }
@@ -429,20 +429,20 @@ impl TxBuilder {
         })
     }
 
-    //     /// Sets the address to drain excess coins to. Usually, when there are excess coins they are sent to a change address
-    //     /// generated by the wallet. This option replaces the usual change address with an arbitrary ScriptPubKey of your choosing.
-    //     /// Just as with a change output, if the drain output is not needed (the excess coins are too small) it will not be included
-    //     /// in the resulting transaction. The only difference is that it is valid to use drain_to without setting any ordinary recipients
-    //     /// with add_recipient (but it is perfectly fine to add recipients as well). If you choose not to set any recipients, you should
-    //     /// either provide the utxos that the transaction should spend via add_utxos, or set drain_wallet to spend all of them.
-    //     /// When bumping the fees of a transaction made with this option, you probably want to use BumpFeeTxBuilder.allow_shrinking
-    //     /// to allow this output to be reduced to pay for the extra fees.
-    //     pub(crate) fn drain_to(&self, script: Arc<Script>) -> Arc<Self> {
-    //         Arc::new(TxBuilder {
-    //             drain_to: Some(script.inner.clone()),
-    //             ..self.clone()
-    //         })
-    //     }
+    /// Sets the address to drain excess coins to. Usually, when there are excess coins they are sent to a change address
+    /// generated by the wallet. This option replaces the usual change address with an arbitrary ScriptPubKey of your choosing.
+    /// Just as with a change output, if the drain output is not needed (the excess coins are too small) it will not be included
+    /// in the resulting transaction. The only difference is that it is valid to use drain_to without setting any ordinary recipients
+    /// with add_recipient (but it is perfectly fine to add recipients as well). If you choose not to set any recipients, you should
+    /// either provide the utxos that the transaction should spend via add_utxos, or set drain_wallet to spend all of them.
+    /// When bumping the fees of a transaction made with this option, you probably want to use BumpFeeTxBuilder.allow_shrinking
+    /// to allow this output to be reduced to pay for the extra fees.
+    pub(crate) fn drain_to(&self, script: Arc<Script>) -> Arc<Self> {
+        Arc::new(TxBuilder {
+            drain_to: Some(script.0.clone()),
+            ..self.clone()
+        })
+    }
     //
     //     /// Enable signaling RBF. This will use the default `nsequence` value of `0xFFFFFFFD`.
     //     pub(crate) fn enable_rbf(&self) -> Arc<Self> {
@@ -504,9 +504,9 @@ impl TxBuilder {
         if self.drain_wallet {
             tx_builder.drain_wallet();
         }
-        // if let Some(script) = &self.drain_to {
-        //     tx_builder.drain_to(script.clone());
-        // }
+        if let Some(script) = &self.drain_to {
+            tx_builder.drain_to(script.clone());
+        }
         // if let Some(rbf) = &self.rbf {
         //     match *rbf {
         //         RbfValue::Default => {