From 1d18764bf1a8e4f0d101d0afc1f494a286649aaf Mon Sep 17 00:00:00 2001
From: Sung Yoon Whang <sungyoon@uber.com>
Date: Tue, 21 Nov 2023 13:23:59 -0800
Subject: [PATCH] Clean up main function

Use switch instead of nested if-elses to enhance readability.
---
 README.md          |  6 +++---
 mockgen/archive.go |  5 +++--
 mockgen/mockgen.go | 49 ++++++++++++++++++++++++++++------------------
 3 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/README.md b/README.md
index 6f05a12..422c0ca 100644
--- a/README.md
+++ b/README.md
@@ -33,9 +33,9 @@ go install go.uber.org/mock/mockgen@latest
 ### Archive mode
 
 Archive mode generates mock interfaces from a package archive
-file (.a). It is enabled by using the -archive flag and two
-non-flag arguments: an import path, and a comma-separated
-list of symbols.
+file (.a). It is enabled by using the -archive flag. An import
+path and a comma-separated list of symbols should be provided
+as a non-flag argument to the command.
 
 Example:
 
diff --git a/mockgen/archive.go b/mockgen/archive.go
index 1f79e4d..54333d5 100644
--- a/mockgen/archive.go
+++ b/mockgen/archive.go
@@ -31,8 +31,9 @@ func archiveMode(importPath string, symbols []string, archive string) (*model.Pa
 	}
 
 	pkg := &model.Package{
-		Name:    tp.Name(),
-		PkgPath: tp.Path(),
+		Name:       tp.Name(),
+		PkgPath:    tp.Path(),
+		Interfaces: make([]*model.Interface, 0, len(symbols)),
 	}
 	for _, name := range symbols {
 		m := tp.Scope().Lookup(name)
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index 93db24b..d17fea3 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -83,31 +83,35 @@ func main() {
 	var pkg *model.Package
 	var err error
 	var packageName string
-	if *source != "" {
+
+	// Switch between modes
+	switch {
+	case *source != "": // source mode
 		pkg, err = sourceMode(*source)
-	} else {
-		if flag.NArg() != 2 {
-			usage()
-			log.Fatal("Expected exactly two arguments")
-		}
+	case *archive != "": // archive mode
+		checkArgs()
 		packageName = flag.Arg(0)
 		interfaces := strings.Split(flag.Arg(1), ",")
-		if *archive != "" {
-			pkg, err = archiveMode(packageName, interfaces, *archive)
-		} else {
-			if packageName == "." {
-				dir, err := os.Getwd()
-				if err != nil {
-					log.Fatalf("Get current directory failed: %v", err)
-				}
-				packageName, err = packageNameOfDir(dir)
-				if err != nil {
-					log.Fatalf("Parse package name failed: %v", err)
-				}
+		pkg, err = archiveMode(packageName, interfaces, *archive)
+
+	default: // reflect mode
+		checkArgs()
+		packageName = flag.Arg(0)
+		interfaces := strings.Split(flag.Arg(1), ",")
+
+		if packageName == "." {
+			dir, err := os.Getwd()
+			if err != nil {
+				log.Fatalf("Get current directory failed: %v", err)
+			}
+			packageName, err = packageNameOfDir(dir)
+			if err != nil {
+				log.Fatalf("Parse package name failed: %v", err)
 			}
-			pkg, err = reflectMode(packageName, interfaces)
 		}
+		pkg, err = reflectMode(packageName, interfaces)
 	}
+
 	if err != nil {
 		log.Fatalf("Loading input failed: %v", err)
 	}
@@ -225,6 +229,13 @@ func parseExcludeInterfaces(names string) map[string]struct{} {
 	return namesSet
 }
 
+func checkArgs() {
+	if flag.NArg() != 2 {
+		usage()
+		log.Fatal("Expected exactly two arguments")
+	}
+}
+
 func usage() {
 	_, _ = io.WriteString(os.Stderr, usageText)
 	flag.PrintDefaults()