gotip: add support for downloading branches

Fixes golang/go#39650

Change-Id: I7d6556db36bd6f7ff292a49cac7a59e25f923225
Reviewed-on: https://go-review.googlesource.com/c/dl/+/301910
Reviewed-by: Katie Hockman <katie@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@golang.org>
Trust: Filippo Valsorda <filippo@golang.org>
diff --git a/gotip/main.go b/gotip/main.go
index 7e636ac..8d9c4b1 100644
--- a/gotip/main.go
+++ b/gotip/main.go
@@ -11,7 +11,8 @@
 //
 // And then use the gotip command as if it were your normal go command.
 //
-// To update, run "gotip download" again.
+// To update, run "gotip download" again. This will always download the main branch.
+// To download an alternative branch, run "gotip download BRANCH".
 // To download a specific CL, run "gotip download NUMBER".
 package main
 
diff --git a/internal/version/gotip.go b/internal/version/gotip.go
index 14ce7ed..61515a1 100644
--- a/internal/version/gotip.go
+++ b/internal/version/gotip.go
@@ -32,14 +32,11 @@
 				log.Fatalf("gotip: %v", err)
 			}
 		case 3:
-			if _, err := strconv.Atoi(os.Args[2]); err != nil {
-				log.Fatalf("gotip: invalid CL number: %q", os.Args[2])
-			}
 			if err := installTip(root, os.Args[2]); err != nil {
 				log.Fatalf("gotip: %v", err)
 			}
 		default:
-			log.Fatalf("gotip: usage: gotip download [CL number]")
+			log.Fatalf("gotip: usage: gotip download [CL number | branch name]")
 		}
 		log.Printf("Success. You may now run 'gotip'!")
 		os.Exit(0)
@@ -53,7 +50,7 @@
 	runGo(root)
 }
 
-func installTip(root, clNumber string) error {
+func installTip(root, target string) error {
 	git := func(args ...string) error {
 		cmd := exec.Command("git", args...)
 		cmd.Stdin = os.Stdin
@@ -77,8 +74,10 @@
 		}
 	}
 
-	if clNumber != "" {
-		fmt.Fprintf(os.Stderr, "This will download and execute code from golang.org/cl/%s, continue? [y/n] ", clNumber)
+	// If the argument is a simple decimal number, consider it a CL number.
+	// Otherwise, consider it a branch name. If it's missing, fetch master.
+	if n, _ := strconv.Atoi(target); n >= 1 && strconv.Itoa(n) == target {
+		fmt.Fprintf(os.Stderr, "This will download and execute code from golang.org/cl/%s, continue? [y/n] ", target)
 		var answer string
 		if fmt.Scanln(&answer); answer != "y" {
 			return fmt.Errorf("interrupted")
@@ -93,10 +92,10 @@
 		if err != nil {
 			return fmt.Errorf("failed to list remotes: %v", err)
 		}
-		r := regexp.MustCompile(`refs/changes/\d\d/` + clNumber + `/(\d+)`)
+		r := regexp.MustCompile(`refs/changes/\d\d/` + target + `/(\d+)`)
 		match := r.FindAllStringSubmatch(string(refs), -1)
 		if match == nil {
-			return fmt.Errorf("CL %v not found", clNumber)
+			return fmt.Errorf("CL %v not found", target)
 		}
 		var ref string
 		var patchSet int
@@ -107,7 +106,13 @@
 				ref = m[0]
 			}
 		}
-		log.Printf("Fetching CL %v, Patch Set %v...", clNumber, patchSet)
+		log.Printf("Fetching CL %v, Patch Set %v...", target, patchSet)
+		if err := git("fetch", "origin", ref); err != nil {
+			return fmt.Errorf("failed to fetch %s: %v", ref, err)
+		}
+	} else if target != "" {
+		log.Printf("Fetching branch %v...", target)
+		ref := "refs/heads/" + target
 		if err := git("fetch", "origin", ref); err != nil {
 			return fmt.Errorf("failed to fetch %s: %v", ref, err)
 		}