internal: Generalize topological sort (#520)

Extract the topological sort logic used in branch_delete
into a separate package so that we can reuse it in #508.

A small user-facing change of this is that 'branch delete'
will delete branches in a more deterministic order.

.changes/unreleased/Changed-20241220-091825.yaml

@@ -0,0 +1,3 @@
+kind: Changed
+body: 'branch delete: If multiple branches are provided, they will be deleted in a more predictable order.'
+time: 2024-12-20T09:18:25.243952-08:00

branch_delete.go

@@ -4,11 +4,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"slices"
 	"strings"
 
 	"github.com/charmbracelet/log"
 	"go.abhg.dev/gs/internal/git"
+	"go.abhg.dev/gs/internal/graph"
 	"go.abhg.dev/gs/internal/must"
 	"go.abhg.dev/gs/internal/spice"
 	"go.abhg.dev/gs/internal/spice/state"
@@ -162,11 +164,31 @@ func (cmd *branchDeleteCmd) Run(ctx context.Context, log *log.Logger, view ui.Vi
 		}
 	}
 
+	// Branches may have relationships with each other.
+	// Sort them in topological order: [close to trunk, ..., further from trunk].
+	topoBranches := graph.Toposort(slices.Sorted(maps.Keys(branchesToDelete)),
+		func(branch string) (string, bool) {
+			info := branchesToDelete[branch]
+			// Branches affect each other's deletion order
+			// only if they're based on each other.
+			_, ok := branchesToDelete[info.Base]
+			return info.Base, ok
+		})
+
+	// Actual deletion will happen in the reverse of that order,
+	// deleting branches based on other branches first.
+	slices.Reverse(topoBranches)
+	deleteOrder := make([]*branchInfo, len(topoBranches))
+	for i, name := range topoBranches {
+		deleteOrder[i] = branchesToDelete[name]
+	}
+
 	// For each branch under consideration,
 	// if it's a tracked branch, update the upstacks from it
 	// to point to its base, or the next branch downstack
 	// if the base is also being deleted.
-	for branch, info := range branchesToDelete {
+	for _, info := range deleteOrder {
+		branch := info.Name
 		if !info.Tracked {
 			continue
 		}
@@ -215,29 +237,6 @@ func (cmd *branchDeleteCmd) Run(ctx context.Context, log *log.Logger, view ui.Vi
 		return fmt.Errorf("checkout %v: %w", checkoutTarget, err)
 	}
 
-	// Remaining branches may have relationships with each other.
-	// We'll need to delete them in topological order: leaf to root.
-	var (
-		deleteOrder []*branchInfo
-		visit       func(string)
-	)
-	visit = func(branch string) {
-		info := branchesToDelete[branch]
-		if info == nil {
-			return // already visited or not in the list
-		}
-
-		visit(info.Base)
-		deleteOrder = append(deleteOrder, info)
-		delete(branchesToDelete, branch)
-	}
-	for branch := range branchesToDelete {
-		visit(branch)
-	}
-
-	// deleteOrder is now in [base, ..., leaf] order. Reverse it.
-	slices.Reverse(deleteOrder)
-
 	branchTx := store.BeginBranchTx()
 	var untrackedNames []string
 	for _, b := range deleteOrder {

internal/graph/doc.go

@@ -0,0 +1,2 @@
+// Package graph provides general-use graph algorithm implementations.
+package graph

internal/graph/topo.go

@@ -0,0 +1,39 @@
+package graph
+
+import "go.abhg.dev/gs/internal/must"
+
+// Toposort performs a topological sort of the given nodes.
+// parent returns the parent of a node, or false if the node doesn't have one.
+//
+// Values returned by parents MUST be in nodes.
+// The graph MUST NOT have a cycle.
+func Toposort[N comparable](
+	nodes []N,
+	parent func(N) (N, bool),
+) []N {
+	topo := make([]N, 0, len(nodes))
+	seen := make(map[N]struct{})
+	var visit func(N)
+	visit = func(n N) {
+		if _, ok := seen[n]; ok {
+			return
+		}
+		seen[n] = struct{}{}
+
+		if p, ok := parent(n); ok {
+			visit(p)
+		}
+
+		topo = append(topo, n)
+	}
+
+	for _, n := range nodes {
+		visit(n)
+	}
+	must.BeEqualf(len(nodes), len(topo),
+		"topological sort produced incorrect number of elements:\n"+
+			"nodes: %v\n"+
+			"topo: %v", nodes, topo)
+
+	return topo
+}

internal/graph/topo_test.go

@@ -0,0 +1,70 @@
+package graph_test
+
+import (
+	"maps"
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.abhg.dev/gs/internal/graph"
+)
+
+func TestToposort(t *testing.T) {
+	tests := []struct {
+		name string
+
+		give map[string][]string // parent -> children
+		want []string
+	}{
+		{name: "Empty", want: []string{}},
+		{
+			name: "Linear",
+			give: map[string][]string{
+				"a": {"b"},
+				"b": {"c"},
+				"c": {"d"},
+			},
+			want: []string{"a", "b", "c", "d"},
+		},
+		{
+			name: "Disjoint",
+			give: map[string][]string{
+				// a -> {b -> d, c}
+				"a": {"b", "c"},
+				"b": {"d"},
+
+				// e -> {f, g}
+				"e": {"f", "g"},
+			},
+			want: []string{
+				"a", "b", "c", "d", "e", "f", "g",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			nodeSet := make(map[string]struct{})
+			parents := make(map[string]string) // node -> parent
+			for parent, children := range tt.give {
+				nodeSet[parent] = struct{}{}
+				for _, child := range children {
+					if p, ok := parents[child]; ok {
+						t.Fatalf("invalid test case: %q already has a parent: %q", child, p)
+					}
+
+					nodeSet[child] = struct{}{}
+					parents[child] = parent
+				}
+			}
+
+			nodes := slices.Sorted(maps.Keys(nodeSet))
+			got := graph.Toposort(nodes, func(n string) (string, bool) {
+				parent, ok := parents[n]
+				return parent, ok
+			})
+
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}