From 81c90f5c73732536c4b27f8139f4299bed06e93a Mon Sep 17 00:00:00 2001
From: MarkSackerberg <93528482+MarkSackerberg@users.noreply.github.com>
Date: Thu, 4 Nov 2021 22:49:23 +0100
Subject: [PATCH] Allow GC to close empty associated accounts (#2554)

* Allow GC to close empty associated accounts

Empty associated accounts will only be closed in case the new flag del_associated_accounts is set. Otherwise behaviour is as before.

Useage:
spl-token gc --del_associated_accounts

* use kebab-case

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>

* align coding style

Thank you jon!

* Add sanity check

* correcting amount check before closing

* formatting correcly

using cargo fmt

* changing sanity check to assert!

Co-authored-by: Sack <none>
Co-authored-by: Jon Cinque <jon.cinque@gmail.com>
Co-authored-by: Mark Sackerberg <Contact@Discord>
---
 token/cli/src/main.rs | 39 +++++++++++++++++++++++++++++++++------
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/token/cli/src/main.rs b/token/cli/src/main.rs
index 4b2260003cb..0d20cbbe9dd 100644
--- a/token/cli/src/main.rs
+++ b/token/cli/src/main.rs
@@ -1214,7 +1214,11 @@ fn command_multisig(config: &Config, address: Pubkey) -> CommandResult {
     Ok(None)
 }
 
-fn command_gc(config: &Config, owner: Pubkey) -> CommandResult {
+fn command_gc(
+    config: &Config,
+    owner: Pubkey,
+    close_empty_associated_accounts: bool,
+) -> CommandResult {
     println_display(config, "Fetching token accounts".to_string());
     let accounts = config
         .rpc_client
@@ -1295,9 +1299,18 @@ fn command_gc(config: &Config, owner: Pubkey) -> CommandResult {
         }
 
         for (address, (amount, decimals, frozen, close_authority)) in accounts {
-            if address == associated_token_account {
-                // leave the associated token account alone
-                continue;
+            match (
+                address == associated_token_account,
+                close_empty_associated_accounts,
+                total_balance > 0,
+            ) {
+                (true, _, true) => continue, // don't ever close associated token account with amount
+                (true, false, _) => continue, // don't close associated token account if close_empty_associated_accounts isn't set
+                (true, true, false) => println_display(
+                    config,
+                    format!("Closing Account {}", associated_token_account),
+                ),
+                _ => {}
             }
 
             if frozen {
@@ -1307,8 +1320,12 @@ fn command_gc(config: &Config, owner: Pubkey) -> CommandResult {
 
             let mut account_instructions = vec![];
 
-            // Transfer the account balance into the associated token account
+            // Sanity check!
+            // we shouldn't ever be here, but if we are here, abort!
+            assert!(amount == 0 || address != associated_token_account);
+
             if amount > 0 {
+                // Transfer the account balance into the associated token account
                 account_instructions.push(transfer_checked(
                     &spl_token::id(),
                     &address,
@@ -2128,6 +2145,12 @@ fn main() {
             SubCommand::with_name("gc")
                 .about("Cleanup unnecessary token accounts")
                 .arg(owner_keypair_arg())
+                .arg(
+                    Arg::with_name("close_empty_associated_accounts")
+                    .long("close-empty-associated-accounts")
+                    .takes_value(false)
+                    .help("close all empty associated token accounts (to get SOL back)")
+                )
         )
         .subcommand(
             SubCommand::with_name("sync-native")
@@ -2609,11 +2632,15 @@ fn main() {
                 }
                 _ => {}
             }
+
+            let close_empty_associated_accounts =
+                matches.is_present("close_empty_associated_accounts");
+
             let (owner_signer, owner_address) =
                 config.signer_or_default(arg_matches, "owner", &mut wallet_manager);
             bulk_signers.push(owner_signer);
 
-            command_gc(&config, owner_address)
+            command_gc(&config, owner_address, close_empty_associated_accounts)
         }
         ("sync-native", Some(arg_matches)) => {
             let address = config.associated_token_address_for_token_or_override(