From 64f97bd19ae4833e84f1da721607a01682676985 Mon Sep 17 00:00:00 2001
From: Evan Wallace <evan.exe@gmail.com>
Date: Fri, 21 Jul 2023 22:57:19 -0400
Subject: [PATCH] fix #3266: support `--analyze` with `--serve`

---
 pkg/cli/cli_impl.go | 103 +++++++++++++++++++++++++-------------------
 1 file changed, 58 insertions(+), 45 deletions(-)

diff --git a/pkg/cli/cli_impl.go b/pkg/cli/cli_impl.go
index 80d75dfe622..74652f8cd23 100644
--- a/pkg/cli/cli_impl.go
+++ b/pkg/cli/cli_impl.go
@@ -1077,13 +1077,59 @@ func splitWithEmptyCheck(s string, sep string) []string {
 	return strings.Split(s, sep)
 }
 
-func runImpl(osArgs []string) int {
-	analyze := false
-	analyzeVerbose := false
+type analyzeMode uint8
+
+const (
+	analyzeDisabled analyzeMode = iota
+	analyzeEnabled
+	analyzeVerbose
+)
+
+func filterAnalyzeFlags(osArgs []string) ([]string, analyzeMode) {
+	analyze := analyzeDisabled
 	end := 0
+	for _, arg := range osArgs {
+		switch arg {
+		case "--analyze":
+			analyze = analyzeEnabled
+		case "--analyze=verbose":
+			analyze = analyzeVerbose
+		default:
+			osArgs[end] = arg
+			end++
+		}
+	}
+	return osArgs[:end], analyze
+}
+
+// Print metafile analysis after the build if it's enabled
+func addAnalyzePlugin(buildOptions *api.BuildOptions, analyze analyzeMode, osArgs []string) {
+	buildOptions.Plugins = append(buildOptions.Plugins, api.Plugin{
+		Name: "PrintAnalysis",
+		Setup: func(build api.PluginBuild) {
+			color := logger.OutputOptionsForArgs(osArgs).Color
+			build.OnEnd(func(result *api.BuildResult) (api.OnEndResult, error) {
+				if result.Metafile != "" {
+					logger.PrintTextWithColor(os.Stderr, color, func(colors logger.Colors) string {
+						return api.AnalyzeMetafile(result.Metafile, api.AnalyzeMetafileOptions{
+							Color:   colors != logger.Colors{},
+							Verbose: analyze == analyzeVerbose,
+						})
+					})
+					os.Stderr.WriteString("\n")
+				}
+				return api.OnEndResult{}, nil
+			})
+		},
+	})
 
+	// Always generate a metafile if we're analyzing, even if it won't be written out
+	buildOptions.Metafile = true
+}
+
+func runImpl(osArgs []string) int {
+	// Special-case running a server
 	for _, arg := range osArgs {
-		// Special-case running a server
 		if arg == "--serve" ||
 			strings.HasPrefix(arg, "--serve=") ||
 			strings.HasPrefix(arg, "--servedir=") ||
@@ -1091,25 +1137,13 @@ func runImpl(osArgs []string) int {
 			serveImpl(osArgs)
 			return 1 // There was an error starting the server if we get here
 		}
-
-		// Special-case analyze just for our CLI
-		if arg == "--analyze" {
-			analyze = true
-			analyzeVerbose = false
-			continue
-		}
-		if arg == "--analyze=verbose" {
-			analyze = true
-			analyzeVerbose = true
-			continue
-		}
-
-		osArgs[end] = arg
-		end++
 	}
-	osArgs = osArgs[:end]
 
+	osArgs, analyze := filterAnalyzeFlags(osArgs)
 	buildOptions, transformOptions, extras, err := parseOptionsForRun(osArgs)
+	if analyze != analyzeDisabled {
+		addAnalyzePlugin(buildOptions, analyze, osArgs)
+	}
 
 	switch {
 	case buildOptions != nil:
@@ -1244,36 +1278,11 @@ func runImpl(osArgs []string) int {
 			}
 		}
 
-		// Print metafile analysis after the build if it's enabled
-		var printAnalysis func(metafile string)
-		if analyze {
-			printAnalysis = func(metafile string) {
-				if metafile == "" {
-					return
-				}
-				logger.PrintTextWithColor(os.Stderr, logger.OutputOptionsForArgs(osArgs).Color, func(colors logger.Colors) string {
-					return api.AnalyzeMetafile(metafile, api.AnalyzeMetafileOptions{
-						Color:   colors != logger.Colors{},
-						Verbose: analyzeVerbose,
-					})
-				})
-				os.Stderr.WriteString("\n")
-			}
-
-			// Always generate a metafile if we're analyzing, even if it won't be written out
-			buildOptions.Metafile = true
-		}
-
 		// Handle post-build actions with a plugin so they also work in watch mode
 		buildOptions.Plugins = append(buildOptions.Plugins, api.Plugin{
 			Name: "PostBuildActions",
 			Setup: func(build api.PluginBuild) {
 				build.OnEnd(func(result *api.BuildResult) (api.OnEndResult, error) {
-					// Print our analysis of the metafile
-					if printAnalysis != nil {
-						printAnalysis(result.Metafile)
-					}
-
 					// Write the metafile to the file system
 					if writeMetafile != nil {
 						writeMetafile(result.Metafile)
@@ -1407,11 +1416,15 @@ func serveImpl(osArgs []string) {
 	options.LogLimit = 5
 	options.LogLevel = api.LogLevelInfo
 
+	filteredArgs, analyze := filterAnalyzeFlags(filteredArgs)
 	extras, errWithNote := parseOptionsImpl(filteredArgs, &options, nil, kindInternal)
 	if errWithNote != nil {
 		logger.PrintErrorWithNoteToStderr(osArgs, errWithNote.Text, errWithNote.Note)
 		return
 	}
+	if analyze != analyzeDisabled {
+		addAnalyzePlugin(&options, analyze, osArgs)
+	}
 
 	serveOptions.OnRequest = func(args api.ServeOnRequestArgs) {
 		logger.PrintText(os.Stderr, logger.LevelInfo, filteredArgs, func(colors logger.Colors) string {