Change up tree Writejc/tree-improvements-refactor

5 files changed, 113 insertions, 68 deletions

diff --git a/internal/git/localrepo/tree.go b/internal/git/localrepo/tree.go
index 4b6d7f88e..322f19713 100644
--- a/internal/git/localrepo/tree.go
+++ b/internal/git/localrepo/tree.go
@@ -209,9 +209,9 @@ func (t *treeQueue) pop() *TreeEntry {
 	return e
 }
 
-// walk scans through the output of ls-tree -r, and constructs a TreeEntry
+// populate scans through the output of ls-tree -r, and constructs a TreeEntry
 // object.
-func (t *TreeEntry) walk(
+func (t *TreeEntry) populate(
 	ctx context.Context,
 	repo *Repo,
 ) error {
@@ -271,6 +271,32 @@ func (t *TreeEntry) walk(
 	return nil
 }
 
+func (t *TreeEntry) walk(dirPath string, callback func(path string, entry *TreeEntry) error) error {
+	if t.Type != Tree {
+		return nil
+	}
+
+	if t.Type == Tree {
+		if err := callback(dirPath, t); err != nil {
+			return err
+		}
+	}
+
+	for _, entry := range t.Entries {
+		if entry.Type == Tree {
+			if err := entry.walk(filepath.Join(dirPath, entry.Path), callback); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (t *TreeEntry) Walk(callback func(path string, entry *TreeEntry) error) error {
+	return t.walk(t.Path, callback)
+}
+
 // WriteTree takes a TreeEntry, and does a depth first walk, writing
 // new trees when needed.
 func (repo *Repo) WriteTree(
@@ -357,53 +383,41 @@ func (repo *Repo) writeEntries(ctx context.Context, entries []*TreeEntry) (git.O
 	return treeOID, nil
 }
 
-// getTreeConfig is configuration that can be passed to GetTree.
-type getTreeConfig struct {
+// readTreeConfig is configuration that can be passed to ReadTree.
+type readTreeConfig struct {
 	recursive bool
 	// relativePath is the relative path at which listing of entries should be started.
 	relativePath string
-	walk         bool
 }
 
-// GetTreeOption is an option that modifies the behavior of GetTree()
-type GetTreeOption func(c *getTreeConfig)
+// ReadTreeOption is an option that modifies the behavior of ReadTree()
+type ReadTreeOption func(c *readTreeConfig)
 
 // WithRecursive returns all entries recursively, but "flattened" in the sense
 // that all subtrees and their entries get returned as Entries, each entry with
 // its full path.
-func WithRecursive() GetTreeOption {
-	return func(c *getTreeConfig) {
+func WithRecursive() ReadTreeOption {
+	return func(c *readTreeConfig) {
 		c.recursive = true
 	}
 }
 
-// WithRelativePath will cause GetTree to return a tree at the relative path.
-func WithRelativePath(relativePath string) GetTreeOption {
-	return func(c *getTreeConfig) {
+// WithRelativePath will cause ReadTree to return a tree at the relative path.
+func WithRelativePath(relativePath string) ReadTreeOption {
+	return func(c *readTreeConfig) {
 		c.relativePath = relativePath
 	}
 }
 
-// WithWalk will walk the rest of the tree, filling in each level's Entries.
-func WithWalk() GetTreeOption {
-	return func(c *getTreeConfig) {
-		c.walk = true
-	}
-}
-
-// GetTree gets a tree object with all of the direct children populated.
+// ReadTree gets a tree object with all of the direct children populated.
 // Walk can be called to populate every level of the tree.
-func (repo *Repo) GetTree(ctx context.Context, treeish git.Revision, options ...GetTreeOption) (*TreeEntry, error) {
-	var c getTreeConfig
+func (repo *Repo) ReadTree(ctx context.Context, treeish git.Revision, options ...ReadTreeOption) (*TreeEntry, error) {
+	var c readTreeConfig
 
 	for _, opt := range options {
 		opt(&c)
 	}
 
-	if c.walk && c.recursive {
-		return nil, errors.New("cannot walk and return recursive listings at the same time")
-	}
-
 	var treeOID git.ObjectID
 	var err error
 
@@ -423,21 +437,19 @@ func (repo *Repo) GetTree(ctx context.Context, treeish git.Revision, options ...
 		Mode: "040000",
 	}
 
-	if c.walk {
-		if err := rootEntry.walk(ctx, repo); err != nil {
+	if c.recursive {
+		if err := rootEntry.populate(ctx, repo); err != nil {
 			return nil, err
 		}
-
-		return rootEntry, nil
-	}
-
-	if rootEntry.Entries, err = repo.listEntries(
-		ctx,
-		treeish,
-		c.relativePath,
-		c.recursive,
-	); err != nil {
-		return nil, fmt.Errorf("listEntries: %w", err)
+	} else {
+		if rootEntry.Entries, err = repo.listEntries(
+			ctx,
+			treeish,
+			c.relativePath,
+			c.recursive,
+		); err != nil {
+			return nil, fmt.Errorf("listEntries: %w", err)
+		}
 	}
 
 	return rootEntry, nil
diff --git a/internal/git/localrepo/tree_test.go b/internal/git/localrepo/tree_test.go
index dd4a23207..049b04cd8 100644
--- a/internal/git/localrepo/tree_test.go
+++ b/internal/git/localrepo/tree_test.go
@@ -380,7 +380,7 @@ func TestTreeEntryByPath(t *testing.T) {
 	}
 }
 
-func TestGetTree(t *testing.T) {
+func TestReadTree(t *testing.T) {
 	t.Parallel()
 
 	cfg := testcfg.Build(t)
@@ -414,8 +414,8 @@ func TestGetTree(t *testing.T) {
 	for _, tc := range []struct {
 		desc            string
 		treeish         git.Revision
-		cfg             *getTreeConfig
-		expectedResults []*TreeEntry
+		cfg             *readTreeConfig
+		expectedResults []TreeEntry
 		expectedErr     error
 	}{
 		{
@@ -425,7 +425,7 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "tree with blob",
 			treeish: treeWithBlob.Revision(),
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "100755", Type: Blob, OID: blobID, Path: "executable"},
 				{Mode: "100644", Type: Blob, OID: blobID, Path: "nonexecutable"},
 			},
@@ -433,24 +433,24 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "tree with subtree",
 			treeish: treeWithSubtree.Revision(),
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "040000", Type: Tree, OID: emptyTreeID, Path: "subdir"},
 			},
 		},
 		{
 			desc:    "nested trees",
 			treeish: treeWithNestedSubtrees.Revision(),
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "040000", Type: Tree, OID: treeWithSubtree, Path: "nested-subdir"},
 			},
 		},
 		{
 			desc:    "recursive nested trees",
 			treeish: treeWithNestedSubtrees.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				recursive: true,
 			},
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "040000", Type: Tree, OID: treeWithSubtree, Path: "nested-subdir"},
 				{Mode: "040000", Type: Tree, OID: emptyTreeID, Path: "nested-subdir/subdir"},
 			},
@@ -458,21 +458,21 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "nested subtree",
 			treeish: treeWithNestedSubtrees.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				relativePath: "nested-subdir",
 			},
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "040000", Type: Tree, OID: emptyTreeID, Path: "subdir"},
 			},
 		},
 		{
 			desc:    "nested recursive subtree",
 			treeish: treeWithSubtreeContainingBlob.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				relativePath: "subdir",
 				recursive:    true,
 			},
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "100644", Type: Blob, OID: blobID, Path: "blob"},
 				{Mode: "040000", Type: Tree, OID: treeWithSubtree, Path: "subdir"},
 				{Mode: "040000", Type: Tree, OID: emptyTreeID, Path: "subdir/subdir"},
@@ -481,10 +481,10 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "recursive nested trees and blobs",
 			treeish: treeWithSubtreeAndBlob.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				recursive: true,
 			},
-			expectedResults: []*TreeEntry{
+			expectedResults: []TreeEntry{
 				{Mode: "100644", Type: Blob, OID: blobID, Path: "blob"},
 				{Mode: "040000", Type: Tree, OID: treeWithSubtree, Path: "subdir"},
 				{Mode: "040000", Type: Tree, OID: emptyTreeID, Path: "subdir/subdir"},
@@ -501,7 +501,7 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "valid revision with invalid path",
 			treeish: treeWithSubtree.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				relativePath: "does-not-exist",
 			},
 			expectedErr: git.ErrReferenceNotFound,
@@ -509,7 +509,7 @@ func TestGetTree(t *testing.T) {
 		{
 			desc:    "valid revision with path pointing to blob",
 			treeish: treeWithSubtreeAndBlob.Revision(),
-			cfg: &getTreeConfig{
+			cfg: &readTreeConfig{
 				relativePath: "blob",
 			},
 			expectedErr: ErrNotTreeish,
@@ -521,20 +521,30 @@ func TestGetTree(t *testing.T) {
 		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
-			var options []GetTreeOption
+			var options []ReadTreeOption
 			if tc.cfg != nil && tc.cfg.recursive {
 				options = append(options, WithRecursive())
 			}
 			if tc.cfg != nil && tc.cfg.relativePath != "" {
 				options = append(options, WithRelativePath(tc.cfg.relativePath))
 			}
-			tree, err := repo.GetTree(ctx, tc.treeish, options...)
+			tree, err := repo.ReadTree(ctx, tc.treeish, options...)
 			if tc.expectedErr != nil {
 				require.ErrorIs(t, err, tc.expectedErr)
 				return
 			}
+			var results []TreeEntry
+
+			require.NoError(t, tree.Walk(func(path string, tree *TreeEntry) error {
+				for _, child := range tree.Entries {
+					child.Path = filepath.Join(path, child.Path)
+					results = append(results, *child)
+					results[len(results)-1].Entries = nil
+				}
+
+				return nil
+			}))
 
-			results := tree.Entries
 			require.Equal(t, tc.expectedResults, results)
 		})
 	}
@@ -713,7 +723,7 @@ func requireTreeEquals(
 	treeOID git.ObjectID,
 	expect *TreeEntry,
 ) {
-	tree, err := repo.GetTree(ctx, git.Revision(treeOID))
+	tree, err := repo.ReadTree(ctx, git.Revision(treeOID))
 	require.NoError(t, err)
 
 	for i, entry := range tree.Entries {
@@ -738,7 +748,7 @@ func requireTreeEquals(
 	}
 }
 
-func TestGetTree_WithWalk(t *testing.T) {
+func TestReadTree_WithRecursive(t *testing.T) {
 	cfg := testcfg.Build(t)
 	ctx := testhelper.Context(t)
 
@@ -872,10 +882,10 @@ func TestGetTree_WithWalk(t *testing.T) {
 
 			commitID, expectedTree := tc.setupTree(t, repoPath)
 
-			tree, err := repo.GetTree(
+			tree, err := repo.ReadTree(
 				ctx,
 				git.Revision(commitID),
-				WithWalk(),
+				WithRecursive(),
 			)
 			require.NoError(t, err)
 			require.Equal(t, expectedTree, *tree)
diff --git a/internal/gitaly/service/commit/tree_entries.go b/internal/gitaly/service/commit/tree_entries.go
index 47d267b6c..8d2b8bb3e 100644
--- a/internal/gitaly/service/commit/tree_entries.go
+++ b/internal/gitaly/service/commit/tree_entries.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"path/filepath"
 	"sort"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
@@ -106,7 +107,7 @@ func (s *server) sendTreeEntries(
 			return err
 		}
 
-		tree, err := repo.GetTree(
+		tree, err := repo.ReadTree(
 			ctx, git.Revision(revision),
 			localrepo.WithRecursive(),
 			localrepo.WithRelativePath(path),
@@ -123,10 +124,21 @@ func (s *server) sendTreeEntries(
 				return nil
 			}
 
-			return fmt.Errorf("listing tree entries: %w", err)
+			return fmt.Errorf("reading tree: %w", err)
 		}
 
-		treeEntries := tree.Entries
+		var treeEntries []*localrepo.TreeEntry
+
+		if err := tree.Walk(func(path string, entry *localrepo.TreeEntry) error {
+			for _, child := range entry.Entries {
+				child.Path = filepath.Join(path, child.Path)
+				treeEntries = append(treeEntries, child)
+			}
+
+			return nil
+		}); err != nil {
+			return fmt.Errorf("listing tree entries: %w", err)
+		}
 
 		entries = make([]*gitalypb.TreeEntry, 0, len(treeEntries))
 		for _, entry := range treeEntries {
diff --git a/internal/gitaly/service/operations/submodules.go b/internal/gitaly/service/operations/submodules.go
index a254397f6..d1358703a 100644
--- a/internal/gitaly/service/operations/submodules.go
+++ b/internal/gitaly/service/operations/submodules.go
@@ -71,10 +71,11 @@ func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.
 	// tree with the new tree abcabc. Continue iterating up the tree,
 	// writing a new tree object each time.
 	for {
-		tree, err := quarantineRepo.GetTree(
+		tree, err := quarantineRepo.ReadTree(
 			ctx,
 			git.Revision("refs/heads/"+string(req.GetBranch())),
 			localrepo.WithRelativePath(path),
+			localrepo.WithRecursive(),
 		)
 		if err != nil {
 			if errors.Is(err, git.ErrReferenceNotFound) {
@@ -84,7 +85,17 @@ func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.
 			return "", fmt.Errorf("error reading tree: %w", err)
 		}
 
-		entries := tree.Entries
+		var entries []*localrepo.TreeEntry
+
+		if err := tree.Walk(func(path string, entry *localrepo.TreeEntry) error {
+			for _, child := range entry.Entries {
+				entries = append(entries, child)
+			}
+
+			return nil
+		}); err != nil {
+			return "", fmt.Errorf("listing tree entries: %w", err)
+		}
 
 		var newEntries []*localrepo.TreeEntry
 		var newTreeID git.ObjectID
diff --git a/internal/gitaly/service/operations/submodules_test.go b/internal/gitaly/service/operations/submodules_test.go
index c35ff50f0..bdadcaaf2 100644
--- a/internal/gitaly/service/operations/submodules_test.go
+++ b/internal/gitaly/service/operations/submodules_test.go
@@ -99,7 +99,7 @@ func TestUserUpdateSubmodule(t *testing.T) {
 			},
 		},
 		{
-			desc:    "successful + nested folder",
+			desc:    "successful nested folder",
 			subPath: "foo/sub",
 			branch:  "master",
 			setup: func(repoPath, subRepoPath string, repoProto, subRepoProto *gitalypb.Repository) setupData {