agent-lib/internal/vendor/select.go

package vendor

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/m3tam3re/agent-lib/internal/discovery"
	"github.com/m3tam3re/agent-lib/internal/gitsource"
	"github.com/m3tam3re/agent-lib/internal/lockfile"
)

// SelectOptions parameterize vendor select.
type SelectOptions struct {
	Add     []string
	Remove  []string
	ModeAll bool
}

// Select changes what a vendored source contributes — without advancing the
// pinned revision (that is vendor update's job). Ids are validated against
// the pinned upstream, collisions re-run, and the external area is rewritten
// deterministically to match the new selection.
func Select(workDir, name string, opts SelectOptions, stdout io.Writer) (*ChangeReport, error) {
	lockPath := filepath.Join(workDir, lockfile.FileName)
	lf, err := lockfile.Load(lockPath)
	if err != nil {
		return nil, err
	}
	src, ok := lf.Sources[name]
	if !ok {
		return nil, fmt.Errorf("source %q is not in the lockfile", name)
	}

	newSel, err := planSelection(src, opts, name)
	if err != nil {
		return nil, err
	}

	repo, err := gitsource.Clone(src.URL, src.Ref)
	if err != nil {
		return nil, err
	}
	pinned, err := repo.Pin(src.Rev)
	if err != nil {
		return nil, fmt.Errorf("%w — run vendor update to re-pin", err)
	}
	tree, err := pinned.Tree()
	if err != nil {
		return nil, err
	}
	items, err := discovery.Scan(tree, discovery.FromLockfileDiscovery(derefDiscovery(src)))
	if err != nil {
		return nil, err
	}

	if err := validateSelectIDs(items, src, opts, name); err != nil {
		return nil, err
	}

	selected := applySelection(items, newSel)
	if err := checkCollisions(workDir, lf, name, selected, src.Renames); err != nil {
		return nil, err
	}

	report := &ChangeReport{Source: name, FromRev: src.Rev, ToRev: src.Rev}
	report.Added, report.Removed = membershipDiff(inventorySet(src.Inventory), idSet(selected))

	cleanupStaleStaging(workDir)
	staging := filepath.Join(workDir, externalDir, ".staging-"+name)
	if err := os.RemoveAll(staging); err != nil {
		return nil, err
	}
	if err := materializeTree(tree, selected, src.Renames, staging); err != nil {
		os.RemoveAll(staging)
		return nil, err
	}
	if err := fillChanged(report, name, lf, staging); err != nil {
		os.RemoveAll(staging)
		return nil, err
	}

	oldDir := filepath.Join(workDir, externalDir, name)
	if err := os.RemoveAll(oldDir); err != nil {
		return nil, fmt.Errorf("replacing %s: %w", oldDir, err)
	}
	if err := os.Rename(staging, oldDir); err != nil {
		return nil, fmt.Errorf("swapping in new external area: %w", err)
	}

	src.Selection = newSel
	src.Inventory = inventoryOf(selected)
	src.Warnings = dedupe(warningsOf(items))
	if err := lf.Save(lockPath); err != nil {
		return nil, err
	}

	fmt.Fprintf(stdout, "source %s: selection updated at pinned rev %s (mode %s)\n", name, short(src.Rev), newSel.Mode)
	printChangeReport(stdout, report)
	return report, nil
}

// planSelection derives the new selection from the flags — pure, before any
// upstream contact: --mode all resets to everything; include mode grows and
// shrinks the include list; all mode --remove moves ids onto the exclude
// list.
func planSelection(src *lockfile.Source, opts SelectOptions, name string) (lockfile.Selection, error) {
	if opts.ModeAll {
		if len(opts.Add) > 0 || len(opts.Remove) > 0 {
			return lockfile.Selection{}, fmt.Errorf("--mode all cannot be combined with --add or --remove")
		}
		return lockfile.Selection{Mode: lockfile.ModeAll}, nil
	}
	if len(opts.Add) == 0 && len(opts.Remove) == 0 {
		return lockfile.Selection{}, fmt.Errorf("nothing to do: pass --add <ids>, --remove <ids> or --mode all")
	}
	if src.Selection.Mode == lockfile.ModeAll {
		if len(opts.Add) > 0 {
			return lockfile.Selection{}, fmt.Errorf("source %q: everything is already selected in all mode; use --remove <ids> to exclude items", name)
		}
		return lockfile.Selection{
			Mode:    lockfile.ModeAll,
			Exclude: unionIDs(src.Selection.Exclude, opts.Remove),
		}, nil
	}
	return lockfile.Selection{
		Mode:    lockfile.ModeInclude,
		Include: unionIDs(minusIDs(src.Selection.Include, opts.Remove), opts.Add),
	}, nil
}

// validateSelectIDs checks flag ids against the pinned upstream (and, for
// --remove in include mode, against the current selection) so typos fail
// before anything is mutated.
func validateSelectIDs(items []discovery.Item, src *lockfile.Source, opts SelectOptions, name string) error {
	if len(opts.Add) > 0 {
		if err := checkIncludeEntries(items, opts.Add, name); err != nil {
			return fmt.Errorf("--add: %w", err)
		}
	}
	if len(opts.Remove) > 0 && src.Selection.Mode == lockfile.ModeInclude {
		selected := map[string]bool{}
		for _, id := range src.Selection.Include {
			selected[id] = true
		}
		for _, id := range opts.Remove {
			if !selected[id] {
				return fmt.Errorf("source %q: --remove id %q is not in the current selection (selected: %s)", name, id, strings.Join(sortedKeys(selected), ", "))
			}
		}
	}
	if len(opts.Remove) > 0 && src.Selection.Mode == lockfile.ModeAll {
		if err := checkIncludeEntries(items, opts.Remove, name); err != nil {
			return fmt.Errorf("--remove: %w", err)
		}
	}
	return nil
}

func unionIDs(a, b []string) []string {
	seen := map[string]bool{}
	for _, id := range a {
		seen[id] = true
	}
	out := append([]string{}, a...)
	for _, id := range b {
		if !seen[id] {
			out = append(out, id)
			seen[id] = true
		}
	}
	sort.Strings(out)
	return out
}

func minusIDs(a, remove []string) []string {
	drop := map[string]bool{}
	for _, id := range remove {
		drop[id] = true
	}
	var out []string
	for _, id := range a {
		if !drop[id] {
			out = append(out, id)
		}
	}
	sort.Strings(out)
	return out
}