From f1f260eb598745ebc59df542980b0a0c8c1b8d06 Mon Sep 17 00:00:00 2001
From: maxlandon <maximelandon@gmail.com>
Date: Tue, 8 Aug 2023 04:14:47 +0200
Subject: [PATCH] Don't use global map for flag completions.

---
 bash_completions.go |  7 ++++---
 command.go          |  8 ++++++++
 completions.go      | 24 +++++++++++-------------
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/bash_completions.go b/bash_completions.go
index 8a531518..091aaff1 100644
--- a/bash_completions.go
+++ b/bash_completions.go
@@ -534,9 +534,9 @@ func writeLocalNonPersistentFlag(buf io.StringWriter, flag *pflag.Flag) {
 
 // prepareCustomAnnotationsForFlags setup annotations for go completions for registered flags
 func prepareCustomAnnotationsForFlags(cmd *Command) {
-	flagCompletionMutex.RLock()
-	defer flagCompletionMutex.RUnlock()
-	for flag := range flagCompletionFunctions {
+	cmd.flagCompletionMutex.RLock()
+	defer cmd.flagCompletionMutex.RUnlock()
+	for flag := range cmd.flagCompletionFunctions {
 		// Make sure the completion script calls the __*_go_custom_completion function for
 		// every registered flag.  We need to do this here (and not when the flag was registered
 		// for completion) so that we can know the root command name for the prefix
@@ -644,6 +644,7 @@ func writeCmdAliases(buf io.StringWriter, cmd *Command) {
 	WriteStringAndCheck(buf, `    fi`)
 	WriteStringAndCheck(buf, "\n")
 }
+
 func writeArgAliases(buf io.StringWriter, cmd *Command) {
 	WriteStringAndCheck(buf, "    noun_aliases=()\n")
 	sort.Strings(cmd.ArgAliases)
diff --git a/command.go b/command.go
index 01f7c6f1..93a470c8 100644
--- a/command.go
+++ b/command.go
@@ -26,6 +26,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
 
 	flag "github.com/spf13/pflag"
 )
@@ -158,6 +159,13 @@ type Command struct {
 	// that we can use on every pflag set and children commands
 	globNormFunc func(f *flag.FlagSet, name string) flag.NormalizedName
 
+	// flagsCompletions contrains completions for arbitrary lists of flags.
+	// Those flags may or may not actually strictly belong to the command in the function,
+	// but registering completions for them through the command allows for garbage-collecting.
+	flagCompletionFunctions map[*flag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
+	// lock for reading and writing from flagCompletionFunctions
+	flagCompletionMutex *sync.RWMutex
+
 	// usageFunc is usage func defined by user.
 	usageFunc func(*Command) error
 	// usageTemplate is usage template defined by user.
diff --git a/completions.go b/completions.go
index ee38c4d0..27110e07 100644
--- a/completions.go
+++ b/completions.go
@@ -18,7 +18,6 @@ import (
 	"fmt"
 	"os"
 	"strings"
-	"sync"
 
 	"github.com/spf13/pflag"
 )
@@ -33,10 +32,10 @@ const (
 )
 
 // Global map of flag completion functions. Make sure to use flagCompletionMutex before you try to read and write from it.
-var flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
-
-// lock for reading and writing from flagCompletionFunctions
-var flagCompletionMutex = &sync.RWMutex{}
+// var flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
+//
+// // lock for reading and writing from flagCompletionFunctions
+// var flagCompletionMutex = &sync.RWMutex{}
 
 // ShellCompDirective is a bit map representing the different behaviors the shell
 // can be instructed to have once completions have been provided.
@@ -135,13 +134,13 @@ func (c *Command) RegisterFlagCompletionFunc(flagName string, f func(cmd *Comman
 	if flag == nil {
 		return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' does not exist", flagName)
 	}
-	flagCompletionMutex.Lock()
-	defer flagCompletionMutex.Unlock()
+	c.flagCompletionMutex.Lock()
+	defer c.flagCompletionMutex.Unlock()
 
-	if _, exists := flagCompletionFunctions[flag]; exists {
+	if _, exists := c.flagCompletionFunctions[flag]; exists {
 		return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' already registered", flagName)
 	}
-	flagCompletionFunctions[flag] = f
+	c.flagCompletionFunctions[flag] = f
 	return nil
 }
 
@@ -479,9 +478,9 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 	// Find the completion function for the flag or command
 	var completionFn func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
 	if flag != nil && flagCompletion {
-		flagCompletionMutex.RLock()
-		completionFn = flagCompletionFunctions[flag]
-		flagCompletionMutex.RUnlock()
+		finalCmd.flagCompletionMutex.RLock()
+		completionFn = finalCmd.flagCompletionFunctions[flag]
+		finalCmd.flagCompletionMutex.RUnlock()
 	} else {
 		completionFn = finalCmd.ValidArgsFunction
 	}
@@ -805,7 +804,6 @@ to your powershell profile.
 				return cmd.Root().GenPowerShellCompletion(out)
 			}
 			return cmd.Root().GenPowerShellCompletionWithDesc(out)
-
 		},
 	}
 	if haveNoDescFlag {