Restore puts submodules on a branch at their pin

A restored submodule whose pin is a branch tip ends on that branch instead
of a detached HEAD: the local branch tracking the ref, else the ref created
as a tracking branch, else another local branch, else another remote branch
at the pin. No branch is moved; in-sync submodules are included (D28).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

internal/submodule/git/provider.go

@@ -327,17 +327,10 @@ func pinState(ctx context.Context, in info, st *submodule.Status, pin, head, bra
 		return fmt.Errorf("pinned commit %.7s not fetched", pin)
 	}
 
-	// remote branches containing the pin, oldest tip first
-	out, err := run(ctx, in.dir, "for-each-ref", "--contains", pin, "--sort=committerdate", "--format=%(refname)", "refs/remotes")
+	containing, err := remoteContaining(ctx, in.dir, pin)
 	if err != nil {
 		return err
 	}
-	var containing []string
-	for _, r := range strings.Fields(out) {
-		if !strings.HasSuffix(r, "/HEAD") {
-			containing = append(containing, strings.TrimPrefix(r, "refs/remotes/"))
-		}
-	}
 	st.Pin.Pushed = len(containing) > 0
 
 	var ref string
@@ -359,6 +352,7 @@ func pinState(ctx context.Context, in info, st *submodule.Status, pin, head, bra
 		return err
 	}
 	st.CheckedOut = st.Pinned
+	var out string
 	if head != pin {
 		if st.CheckedOut, err = commit(ctx, in.dir, head); err != nil {
 			return err
@@ -417,6 +411,22 @@ func pinState(ctx context.Context, in info, st *submodule.Status, pin, head, bra
 	return nil
 }
 
+// remoteContaining lists the remote branches containing pin, oldest tip
+// first, without the origin/HEAD symref.
+func remoteContaining(ctx context.Context, dir, pin string) ([]string, error) {
+	out, err := run(ctx, dir, "for-each-ref", "--contains", pin, "--sort=committerdate", "--format=%(refname)", "refs/remotes")
+	if err != nil {
+		return nil, err
+	}
+	var containing []string
+	for _, r := range strings.Fields(out) {
+		if !strings.HasSuffix(r, "/HEAD") {
+			containing = append(containing, strings.TrimPrefix(r, "refs/remotes/"))
+		}
+	}
+	return containing, nil
+}
+
 // resolve finds the ref the pin is compared with, following the order in
 // docs/state-model.md: the target (see targetRef) if the pin is on it, else the
 // branch containing the pin that is closest to it, else the target anyway,

internal/submodule/git/restore.go

@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"io/fs"
 	"os"
+	"path"
 	"path/filepath"
+	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -45,7 +47,7 @@ func (p *Provider) RestoreLoss(ctx context.Context, path string, to submodule.Re
 // submodule is cloned and missing commits are fetched; git never prompts.
 func (p *Provider) Restore(ctx context.Context, path string, to submodule.RestoreTarget) error {
 	if path == "" {
-		return p.restoreRepo(ctx, p.root, "HEAD")
+		return p.restoreRepo(ctx, p.root, "HEAD", nil)
 	}
 	p.mu.RLock()
 	in, ok := p.mods[path]
@@ -67,7 +69,7 @@ func (p *Provider) Restore(ctx context.Context, path string, to submodule.Restor
 			return fmt.Errorf("%s: %w", path, err)
 		}
 	}
-	return p.restoreRepo(ctx, in.dir, pin)
+	return p.restoreRepo(ctx, in.dir, pin, &in)
 }
 
 // restoreTarget returns the directory and the commit path is restored to:
@@ -123,7 +125,9 @@ func pinOf(ctx context.Context, in info, to submodule.RestoreTarget) (string, er
 // or "HEAD") and then every submodule its index records. Branches are never
 // moved: a checkout that is not on target is detached first. If the checkout
 // stays on target, its staged pins are kept and are the submodules' targets.
-func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
+// For a submodule (in is not nil) a detached HEAD is then put on a branch at
+// target if there is one, see attachBranch; the submodules below always are.
+func (p *Provider) restoreRepo(ctx context.Context, dir, target string, in *info) error {
 	env := p.networkEnv()
 	g := func(args ...string) (string, error) {
 		return p.exec(ctx, dir, env, append(append([]string{}, restoreArgs...), args...)...)
@@ -138,6 +142,7 @@ func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
 	if gitDir, _, err := gitDirs(dir); err == nil {
 		p.clearOperations(ctx, dir, gitDir)
 	}
+	prev, _, _ := try(ctx, dir, "symbolic-ref", "--short", "-q", "HEAD") // for the ref, see attachBranch
 	var keep []gitlink
 	if head, err := g("rev-parse", "--verify", "-q", "HEAD"); err != nil || head != sha {
 		if _, err := g("checkout", "--force", "--detach", sha); err != nil {
@@ -161,6 +166,11 @@ func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
 	if _, err := g("submodule", "sync", "-q"); err != nil {
 		return err
 	}
+	if in != nil {
+		if err := p.attachBranch(ctx, *in, sha, prev); err != nil {
+			return err
+		}
+	}
 
 	links, err := indexGitlinks(ctx, dir)
 	if err != nil {
@@ -178,6 +188,16 @@ func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
 		}
 	}
 
+	// the .gitmodules branch of each child, for its ref
+	branches := map[string]string{}
+	if out, ok, err := try(ctx, dir, "config", "-f", ".gitmodules", "--get-regexp", `^submodule\..*\.(path|branch)$`); err != nil {
+		return err
+	} else if ok {
+		for _, m := range parseGitmodules(out) {
+			branches[path.Clean(m.path)] = m.branch
+		}
+	}
+
 	// the children are independent repositories
 	var wg sync.WaitGroup
 	errs := make([]error, len(links))
@@ -185,7 +205,8 @@ func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			if err := p.restoreRepo(ctx, filepath.Join(dir, filepath.FromSlash(l.path)), l.sha); err != nil {
+			child := info{superDir: dir, dir: filepath.Join(dir, filepath.FromSlash(l.path)), relPath: l.path, branch: branches[l.path]}
+			if err := p.restoreRepo(ctx, child.dir, l.sha, &child); err != nil {
 				errs[i] = err
 			}
 		}()
@@ -194,6 +215,135 @@ func (p *Provider) restoreRepo(ctx context.Context, dir, target string) error {
 	return errors.Join(errs...)
 }
 
+// attachBranch puts the detached HEAD of the restored submodule in on a
+// branch whose tip is sha, the first that git accepts of, in this order:
+//
+//  1. a local branch tracking the ref the UI shows (see resolve), or named
+//     like it;
+//  2. the ref itself, as a new local branch tracking it, if no local branch
+//     has its name;
+//  3. any other local branch, alphabetically;
+//  4. any other remote branch whose name no local branch has, as a new
+//     tracking branch; the ref's remote first, then alphabetically.
+//
+// A HEAD that is on a branch stays there. No branch is moved: a local branch
+// elsewhere is no candidate, and its remote branch neither. If nothing fits,
+// HEAD stays detached; only a context error is returned. prev is the branch
+// checked out before the restore; the ref depends on it.
+func (p *Provider) attachBranch(ctx context.Context, in info, sha, prev string) error {
+	if _, onBranch, err := try(ctx, in.dir, "symbolic-ref", "-q", "HEAD"); err != nil || onBranch {
+		return err
+	}
+	out, ok, err := try(ctx, in.dir, "for-each-ref", "--points-at="+sha,
+		"--format=%(refname)%00%(upstream)%00%(symref)", "refs/heads", "refs/remotes")
+	if err != nil || !ok {
+		return err
+	}
+	var locals, remotes []string // alphabetical, as for-each-ref sorts them
+	upstreams := map[string]string{}
+	for _, line := range strings.Split(out, "\n") {
+		f := strings.Split(line, "\x00")
+		if len(f) != 3 || f[2] != "" { // a symref is origin/HEAD
+			continue
+		}
+		if name, ok := strings.CutPrefix(f[0], "refs/heads/"); ok {
+			locals = append(locals, name)
+			upstreams[name] = strings.TrimPrefix(f[1], "refs/remotes/")
+		} else if name, ok := strings.CutPrefix(f[0], "refs/remotes/"); ok {
+			remotes = append(remotes, name)
+		}
+	}
+	if len(locals)+len(remotes) == 0 {
+		return nil
+	}
+
+	// the ref, if it is a remote branch
+	var ref string
+	containing, err := remoteContaining(ctx, in.dir, sha)
+	if err != nil {
+		return err
+	}
+	if r, source, _, err := resolve(ctx, in, sha, prev, containing); err != nil {
+		return err
+	} else if source != submodule.RefFromTag && source != submodule.RefFromSHA {
+		ref = r
+	}
+
+	localOf := func(remote string) string { _, l, _ := strings.Cut(remote, "/"); return l }
+	hasLocal := func(name string) (bool, error) {
+		_, ok, err := try(ctx, in.dir, "rev-parse", "--verify", "-q", "refs/heads/"+name)
+		return ok, err
+	}
+	type candidate struct{ local, remote string } // remote: create local, tracking it
+	var cands []candidate
+	addRemote := func(r string) error {
+		l := localOf(r)
+		if l == "" {
+			return nil
+		}
+		exists, err := hasLocal(l)
+		if err == nil && !exists {
+			cands = append(cands, candidate{l, r})
+		}
+		return err
+	}
+	if ref != "" {
+		for _, l := range locals {
+			if upstreams[l] == ref {
+				cands = append(cands, candidate{local: l})
+			}
+		}
+		if l := localOf(ref); slices.Contains(locals, l) {
+			cands = append(cands, candidate{local: l})
+		}
+		if slices.Contains(remotes, ref) {
+			if err := addRemote(ref); err != nil {
+				return err
+			}
+		}
+	}
+	for _, l := range locals {
+		cands = append(cands, candidate{local: l})
+	}
+	preferred, _, _ := strings.Cut(ref, "/")
+	if preferred == "" {
+		if preferred, err = remoteName(ctx, in.dir); err != nil {
+			return err
+		}
+	}
+	slices.SortStableFunc(remotes, func(a, b string) int {
+		ap, bp := strings.HasPrefix(a, preferred+"/"), strings.HasPrefix(b, preferred+"/")
+		switch {
+		case ap && !bp:
+			return -1
+		case bp && !ap:
+			return 1
+		}
+		return 0
+	})
+	for _, r := range remotes {
+		if r != ref {
+			if err := addRemote(r); err != nil {
+				return err
+			}
+		}
+	}
+
+	for _, c := range cands {
+		args := append(append([]string{}, restoreArgs...), "switch", "-q", c.local)
+		if c.remote != "" {
+			args = append(append([]string{}, restoreArgs...), "switch", "-q", "-c", c.local, "--track", c.remote)
+		}
+		if _, err := p.exec(ctx, in.dir, nil, args...); err == nil {
+			return nil
+		} else if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		// e.g. checked out in another worktree: try the next one
+	}
+	return nil
+}
+
 // networkEnv is the environment of commands that may clone or fetch.
 func (p *Provider) networkEnv() []string {
 	return append(append([]string{}, p.sshEnv...), p.extraEnv...)

internal/submodule/git/restore_test.go

@@ -332,3 +332,130 @@ func TestRestoreLossCountsStagedPinsOfAMovingRepository(t *testing.T) {
 		t.Errorf("loss %+v, want 1 staged pin and the 2 detached nested commits", l)
 	}
 }
+
+// restoreSub restores the submodule path and returns the branch its HEAD is
+// on ("" if detached).
+func restoreSub(t *testing.T, p *Provider, super, path string) string {
+	t.Helper()
+	if err := p.Restore(context.Background(), path, submodule.RestoreStaged); err != nil {
+		t.Fatal(err)
+	}
+	return out(t, filepath.Join(super, filepath.FromSlash(path)), "branch", "--show-current")
+}
+
+func TestRestoreAttachesTheLocalBranchTrackingTheRef(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	// src/ui: in sync, the pin is behind origin/main, local main is at the pin
+	git(t, filepath.Join(super, "src", "ui"), "checkout", "-q", "--detach")
+	git(t, filepath.Join(super, "src", "ui"), "branch", "aaa") // alphabetically first, but no ref
+
+	if got := restoreSub(t, p, super, "src/ui"); got != "main" {
+		t.Errorf("branch = %q, want main", got)
+	}
+}
+
+func TestRestoreCreatesTheRefBranchBeforeOtherBranches(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	log := filepath.Join(super, "libs", "log")
+	// libs/log: the pin is the tip of origin/main and origin/archive
+	git(t, log, "checkout", "-q", "--detach")
+	git(t, log, "branch", "-D", "main")
+
+	if got := restoreSub(t, p, super, "libs/log"); got != "main" {
+		t.Errorf("branch = %q, want main", got)
+	}
+	if up := out(t, log, "rev-parse", "--abbrev-ref", "main@{upstream}"); up != "origin/main" {
+		t.Errorf("upstream = %q, want origin/main", up)
+	}
+	if out(t, log, "branch", "--list", "archive") != "" {
+		t.Error("archive must not be created")
+	}
+}
+
+func TestRestoreCreatesARemoteFeatureBranch(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	topic := filepath.Join(super, "libs", "topic")
+	mainTip := out(t, topic, "rev-parse", "main")
+	// libs/topic: detached at the tip of origin/feat/search, local main elsewhere
+
+	if got := restoreSub(t, p, super, "libs/topic"); got != "feat/search" {
+		t.Errorf("branch = %q, want feat/search", got)
+	}
+	if got := out(t, topic, "rev-parse", "main"); got != mainTip {
+		t.Errorf("main moved: %s -> %s", mainTip, got)
+	}
+}
+
+func TestRestorePrefersLocalBranchesAlphabetically(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	feature := filepath.Join(super, "libs", "feature")
+	// libs/feature: the pin is on no branch tip
+	git(t, feature, "branch", "zz")
+	git(t, feature, "branch", "aa")
+
+	if got := restoreSub(t, p, super, "libs/feature"); got != "aa" {
+		t.Errorf("branch = %q, want aa", got)
+	}
+}
+
+func TestRestoreNeverMovesALocalBranchToAttach(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	topic := filepath.Join(super, "libs", "topic")
+	git(t, topic, "branch", "feat/search", "origin/main") // same name, elsewhere
+	before := out(t, topic, "rev-parse", "feat/search")
+
+	if got := restoreSub(t, p, super, "libs/topic"); got != "" {
+		t.Errorf("branch = %q, want detached", got)
+	}
+	if after := out(t, topic, "rev-parse", "feat/search"); after != before {
+		t.Errorf("feat/search moved: %s -> %s", before, after)
+	}
+}
+
+func TestRestoreKeepsTheCheckedOutBranch(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	log := filepath.Join(super, "libs", "log")
+	git(t, log, "switch", "-q", "-c", "zz") // at the pin, like main
+
+	if got := restoreSub(t, p, super, "libs/log"); got != "zz" {
+		t.Errorf("branch = %q, want zz kept", got)
+	}
+}
+
+func TestRestoreSkipsABranchOfAnotherWorktree(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	log := filepath.Join(super, "libs", "log")
+	git(t, log, "checkout", "-q", "--detach")
+	git(t, log, "worktree", "add", "-q", filepath.Join(root, "wt"), "main")
+
+	if got := restoreSub(t, p, super, "libs/log"); got == "main" {
+		t.Error("main is checked out in another worktree")
+	}
+}
+
+func TestRestoreSuperprojectAttachesNestedSubmodules(t *testing.T) {
+	p, root := ownFixture(t)
+	super := filepath.Join(root, "super")
+	nested := filepath.Join(super, "libs", "log", "modules", "nested")
+	git(t, nested, "checkout", "-q", "--detach")
+
+	if err := p.Restore(context.Background(), "", submodule.RestoreStaged); err != nil {
+		t.Fatal(err)
+	}
+	if got := out(t, nested, "branch", "--show-current"); got != "main" {
+		t.Errorf("nested branch = %q, want main", got)
+	}
+	if got := out(t, filepath.Join(super, "libs", "topic"), "branch", "--show-current"); got != "feat/search" {
+		t.Errorf("libs/topic branch = %q, want feat/search", got)
+	}
+	if got := out(t, super, "branch", "--show-current"); got != "main" {
+		t.Errorf("superproject branch = %q, want main", got)
+	}
+}

internal/submodule/git/switch.go

@@ -107,7 +107,7 @@ func (p *Provider) Switch(ctx context.Context, path string, b submodule.Branch)
 			return "", err
 		}
 	}
-	if err := p.restoreRepo(ctx, dir, b.SHA); err != nil {
+	if err := p.restoreRepo(ctx, dir, b.SHA, nil); err != nil { // HEAD goes on b below
 		return "", err
 	}