sfdupes/trees.go

package main

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"fmt"
	"io"
	"os"
	"slices"
	"strconv"
	"strings"
)

// treeNode is one directory reconstructed from the record paths.
type treeNode struct {
	path   string
	parent *treeNode
	// entries holds the serialized child entries until the digest is
	// computed from them, and is then dropped.
	entries   []string
	digest    [sha256.Size]byte
	fileCount int64
	totalSize int64
}

// treesPhase names one phase of trees with its place among all of them,
// as "[2/4] read", for the phase's progress display. The list here is
// the one place the phases are counted.
func treesPhase(name string) string {
	return phaseLabel([]string{"open", "read", "group", "write"}, name)
}

// runTrees implements the trees subcommand: it reads every record from
// the database in path order, reconstructs the directory hierarchy from
// the record paths, computes a Merkle-style digest per directory, and
// prints maximal duplicate-tree groups as TSV on stdout. It never
// touches the scanned filesystem; its only I/O is the database, stdout,
// and stderr. It names the database on stderr before it opens it. Any
// database problem, including a missing database, is fatal.
//
// The work runs in four phases, each under a progress display of its
// own on stderr: open (open the database), read (count the records,
// then read them into the hierarchy), group (group the directories by
// digest) and write (write the groups to stdout). Each display ends
// before the next phase starts, and before an error is returned to be
// printed.
func runTrees(ctx context.Context, stdout io.Writer) error {
	dbPath, source, err := databasePath()
	if err != nil {
		return err
	}

	fmt.Fprintf(os.Stderr, "trees: database %s (%s)\n",
		escapePath(dbPath), source)

	// Opening can wait on a busy database for as long as the busy
	// timeout.
	prog := newProgress(progressParams{
		label: treesPhase("open"), total: -1,
		waitingFor: "opening the database",
	})
	db, err := openReportDatabase(ctx, dbPath)

	prog.finish()

	if err != nil {
		return err
	}

	defer func() { _ = db.Close() }()

	prog = newProgress(progressParams{
		label: treesPhase("read"), unit: "records", total: -1,
		waitingFor: countingRecords,
	})
	records, super, allDirs, err := readTree(ctx, db, prog)

	prog.finish()

	if err != nil {
		return fmt.Errorf("database %s: %w", dbPath, err)
	}

	prog = newProgress(progressParams{
		label: treesPhase("group"), unit: "directories",
		total: int64(len(allDirs)),
	})
	dupes := collectTreeGroups(allDirs, super, prog)

	prog.finish()

	prog = newProgress(progressParams{
		label: treesPhase("write"), unit: "groups", total: int64(len(dupes)),
	})
	dupeTrees, reclaimable, err := writeTreeGroups(stdout, dupes, prog)

	prog.finish()

	if err != nil {
		return err
	}

	fmt.Fprintf(os.Stderr,
		"trees: %d records read, %d duplicate tree groups, %d dupe trees, "+
			"%s reclaimable\n",
		records, len(dupes), dupeTrees, humanBytes(reclaimable))

	return nil
}

// readTree gives prog the number of records in db as its total, then
// reads every record in path order into the directory hierarchy,
// counting each on prog. It returns the number of records read, the
// super-root and every directory.
func readTree(ctx context.Context, db *sql.DB,
	prog *progress,
) (int, *treeNode, []*treeNode, error) {
	records := 0
	tree := newTreeBuilder()

	err := loadFileRows(ctx, db, prog.setTotal, func(r scanRec) {
		records++

		prog.increment()
		tree.add(r)
	})
	if err != nil {
		return 0, nil, nil, err
	}

	super, allDirs := tree.finish()

	return records, super, allDirs, nil
}

// writeTreeGroups writes dupes to stdout as TSV, a header line and then
// a row for every member of each group but its first, counting each
// group on prog once its rows are written. stdout may be the terminal
// prog is drawn on, so the lines go through prog's writer. It returns
// the number of rows (dupe trees) and the bytes they could reclaim.
func writeTreeGroups(stdout io.Writer, dupes [][]*treeNode,
	prog *progress,
) (int, int64, error) {
	out := prog.writer(stdout)

	_, err := fmt.Fprintln(out, "first\tdupe\tfiles\tsize")
	if err != nil {
		return 0, 0, fmt.Errorf("write stdout: %w", err)
	}

	dupeTrees := 0

	var reclaimable int64

	for _, g := range dupes {
		first := g[0]
		for _, n := range g[1:] {
			_, err = fmt.Fprintf(out, "%s\t%s\t%d\t%d\n",
				escapePath(first.path), escapePath(n.path),
				first.fileCount, first.totalSize)
			if err != nil {
				return 0, 0, fmt.Errorf("write stdout: %w", err)
			}

			dupeTrees++
			reclaimable += first.totalSize
		}

		prog.increment()
	}

	err = out.Flush()
	if err != nil {
		return 0, 0, fmt.Errorf("write stdout: %w", err)
	}

	return dupeTrees, reclaimable, nil
}

// treeBuilder reconstructs the directory hierarchy from records added
// in path order, under a synthetic super-root. Paths are split on "/";
// for absolute paths the first component is empty, which becomes the
// top-level directory with path "/". In path order all the paths under
// one directory come together, so a directory is complete once a path
// outside it is added: its digest is computed then and its entries are
// dropped. Only the directories holding the latest path keep entries.
type treeBuilder struct {
	super *treeNode
	// open lists the directories holding the latest path, outermost
	// first, starting with the super-root; names[i] is open[i]'s name.
	open  []*treeNode
	names []string
	// dirs lists every completed directory.
	dirs []*treeNode
}

func newTreeBuilder() *treeBuilder {
	super := &treeNode{}

	return &treeBuilder{
		super: super,
		open:  []*treeNode{super},
		names: []string{""},
	}
}

// add adds one record. Each record must come after the previous one in
// path order (byte order); otherwise a completed directory would be
// started again as a second directory with the same path.
func (b *treeBuilder) add(r scanRec) {
	comps := strings.Split(r.path, "/")
	dirNames, name := comps[:len(comps)-1], comps[len(comps)-1]

	// Keep the open directories that hold this path; complete the rest.
	depth := 1
	for depth < len(b.open) && depth <= len(dirNames) &&
		b.names[depth] == dirNames[depth-1] {
		depth++
	}

	b.closeTo(depth)

	for _, c := range dirNames[depth-1:] {
		b.openDir(c)
	}

	dir := b.open[len(b.open)-1]
	dir.entries = append(dir.entries, fileEntry(name, r))
	dir.fileCount++
	dir.totalSize += r.size
}

// openDir starts the directory called name inside the innermost open
// one.
func (b *treeBuilder) openDir(name string) {
	parent := b.open[len(b.open)-1]
	path := parent.path + "/" + name

	// The root directory's path is "/", not empty, and its children's
	// paths start with one slash, not two.
	switch {
	case parent == b.super && name == "":
		path = "/"
	case parent == b.super:
		path = name
	case parent.path == "/":
		path = "/" + name
	}

	b.open = append(b.open, &treeNode{path: path, parent: parent})
	b.names = append(b.names, name)
}

// closeTo completes the open directories after the first n, innermost
// first: each one's digest is computed and entered in its parent along
// with its totals.
func (b *treeBuilder) closeTo(n int) {
	for len(b.open) > n {
		last := len(b.open) - 1
		dir, name := b.open[last], b.names[last]
		b.open, b.names = b.open[:last], b.names[:last]

		dir.computeDigest()

		dir.parent.entries = append(dir.parent.entries,
			"d\x00"+name+"\x00"+string(dir.digest[:]))
		dir.parent.fileCount += dir.fileCount
		dir.parent.totalSize += dir.totalSize
		b.dirs = append(b.dirs, dir)
	}
}

// finish completes every open directory and returns the super-root and
// every directory.
func (b *treeBuilder) finish() (*treeNode, []*treeNode) {
	b.closeTo(1)

	return b.super, b.dirs
}

// fileEntry serializes a file child for its directory's digest: its
// name and its signature (size, head, tail, content); mtime is
// excluded.
func fileEntry(name string, r scanRec) string {
	content := r.content

	// A record without a content hash has unknown content (README
	// "Database"): give it a signature no other file can share, so
	// trees containing it never compare equal. Real hashes are hex, so
	// the NUL-prefixed form cannot collide.
	if content == "" {
		content = "unhashed\x00" + r.path
	}

	return "f\x00" + name + "\x00" + strconv.FormatInt(r.size, 10) +
		"\x00" + r.head + "\x00" + r.tail + "\x00" + content
}

// collectTreeGroups groups directories by digest, counting each one on
// prog as it is grouped, and returns every maximal group with two or
// more members, each group's members sorted by path, groups ordered by
// tree size descending then by first path ascending.
func collectTreeGroups(allDirs []*treeNode, super *treeNode,
	prog *progress,
) [][]*treeNode {
	groups := make(map[[sha256.Size]byte][]*treeNode)
	for _, d := range allDirs {
		groups[d.digest] = append(groups[d.digest], d)

		prog.increment()
	}

	var dupes [][]*treeNode

	for _, g := range groups {
		if len(g) < minGroupSize || suppressed(g, super) {
			continue
		}

		slices.SortFunc(g, func(a, b *treeNode) int {
			return strings.Compare(a.path, b.path)
		})
		dupes = append(dupes, g)
	}

	// Biggest reclaimable space first; ties broken by first path.
	slices.SortFunc(dupes, func(a, b []*treeNode) int {
		if a[0].totalSize != b[0].totalSize {
			if a[0].totalSize > b[0].totalSize {
				return -1
			}

			return 1
		}

		return strings.Compare(a[0].path, b[0].path)
	})

	return dupes
}

// computeDigest sets n's digest and drops its entries. A directory's
// digest is the SHA-256 of its child entries — files serialized with
// name and signature, subdirectories with name and recursive digest —
// sorted byte-lexicographically. Filenames cannot contain NUL or "/",
// so NUL delimiters are unambiguous.
func (n *treeNode) computeDigest() {
	slices.Sort(n.entries)

	h := sha256.New()
	for _, e := range n.entries {
		h.Write([]byte(e))
		h.Write([]byte{0})
	}

	copy(n.digest[:], h.Sum(nil))
	n.entries = nil
}

// suppressed reports whether a duplicate-tree group is non-maximal: its
// members' parents are pairwise distinct real directories that all
// share a single digest, so the group is wholly implied by its parents'
// (or a further ancestor's) group. Groups containing siblings (shared
// parent) or members whose parents differ are always reported.
func suppressed(g []*treeNode, super *treeNode) bool {
	seen := make(map[*treeNode]bool, len(g))

	var parentDigest [sha256.Size]byte

	for i, n := range g {
		p := n.parent
		if p == nil || p == super {
			return false
		}

		if seen[p] {
			return false // siblings: not implied by any parent group
		}

		seen[p] = true
		if i == 0 {
			parentDigest = p.digest
		} else if p.digest != parentDigest {
			return false
		}
	}

	return true
}