localrepo: Add GetFullTree helper

GetFullTree is a helper that gets a tree object with all of the entries
at every level. This will be used to modify the tree using other helpers
to then write the modified tree into the repository.

2 files changed, 221 insertions, 0 deletions

diff --git a/internal/git/localrepo/tree.go b/internal/git/localrepo/tree.go
index 01e0c5aea..b8f73fb21 100644
--- a/internal/git/localrepo/tree.go
+++ b/internal/git/localrepo/tree.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"path/filepath"
 	"sort"
 	"strings"
 
@@ -79,6 +80,8 @@ type TreeEntry struct {
 	Path string
 	// Type is the type of the tree entry.
 	Type ObjectType
+
+	Entries []*TreeEntry
 }
 
 // IsBlob returns whether or not the TreeEntry is a blob.
@@ -142,3 +145,76 @@ func (repo *Repo) WriteTree(ctx context.Context, entries []*TreeEntry) (git.Obje
 
 	return treeOID, nil
 }
+
+func depthByPath(path string) int {
+	return len(strings.Split(path, "/"))
+}
+
+type treeQueue []*TreeEntry
+
+func (t treeQueue) first() *TreeEntry {
+	return t[len(t)-1]
+}
+
+func (t *treeQueue) push(e *TreeEntry) {
+	*t = append(*t, e)
+}
+
+func (t *treeQueue) pop() *TreeEntry {
+	e := t.first()
+
+	*t = (*t)[:len(*t)-1]
+
+	return e
+}
+
+// GetFullTree gets a tree object with all of the Entries populated for every
+// level of the tree.
+func (repo *Repo) GetFullTree(ctx context.Context, treeish git.Revision) (*TreeEntry, error) {
+	treeOID, err := repo.ResolveRevision(ctx, git.Revision(fmt.Sprintf("%s^{tree}", treeish)))
+	if err != nil {
+		return nil, fmt.Errorf("getting revision %w", err)
+	}
+
+	rootEntry := &TreeEntry{
+		OID:  treeOID,
+		Type: Tree,
+		Mode: "040000",
+	}
+
+	entries, err := repo.ListEntries(
+		ctx,
+		git.Revision(treeOID),
+		&ListEntriesConfig{
+			Recursive: true,
+		},
+	)
+
+	queue := treeQueue{rootEntry}
+
+	for _, entry := range entries {
+		if depthByPath(entry.Path) < len(queue) {
+			levelsToPop := len(queue) - depthByPath(entry.Path)
+
+			for i := 0; i < levelsToPop; i++ {
+				queue.pop()
+			}
+		}
+
+		entry.Path = filepath.Base(entry.Path)
+		queue.first().Entries = append(
+			queue.first().Entries,
+			entry,
+		)
+
+		if entry.Type == Tree {
+			queue.push(entry)
+		}
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("listing entries: %w", err)
+	}
+
+	return rootEntry, nil
+}
diff --git a/internal/git/localrepo/tree_test.go b/internal/git/localrepo/tree_test.go
index 4d4b6c0f6..43ccac2c0 100644
--- a/internal/git/localrepo/tree_test.go
+++ b/internal/git/localrepo/tree_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/text"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
@@ -228,3 +229,147 @@ func TestWriteTree(t *testing.T) {
 		})
 	}
 }
+
+func TestGetFullTree(t *testing.T) {
+	cfg := testcfg.Build(t)
+	ctx := testhelper.Context(t)
+
+	testCases := []struct {
+		desc        string
+		setupTree   func(t *testing.T, repoPath string) (git.ObjectID, TreeEntry)
+		expectedErr error
+	}{
+		{
+			desc: "flat tree",
+			setupTree: func(t *testing.T, repoPath string) (git.ObjectID, TreeEntry) {
+				blobA := gittest.WriteBlob(t, cfg, repoPath, []byte("a"))
+				blobB := gittest.WriteBlob(t, cfg, repoPath, []byte("b"))
+
+				treeID := gittest.WriteTree(t, cfg, repoPath, []gittest.TreeEntry{
+					{
+						Mode: "100644",
+						Path: "fileA",
+						OID:  blobA,
+					},
+					{
+						Mode: "100644",
+						Path: "fileB",
+						OID:  blobB,
+					},
+				})
+				return gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("master"),
+						gittest.WithTree(treeID)),
+					TreeEntry{
+						OID:  treeID,
+						Mode: "040000",
+						Type: Tree,
+						Entries: []*TreeEntry{
+							{
+								Mode: "100644",
+								Path: "fileA",
+								OID:  blobA,
+								Type: Blob,
+							},
+							{
+								Mode: "100644",
+								Path: "fileB",
+								OID:  blobB,
+								Type: Blob,
+							},
+						},
+					}
+			},
+		},
+		{
+			desc: "nested tree",
+			setupTree: func(t *testing.T, repoPath string) (git.ObjectID, TreeEntry) {
+				blobA := gittest.WriteBlob(t, cfg, repoPath, []byte("a"))
+				blobB := gittest.WriteBlob(t, cfg, repoPath, []byte("b"))
+				blobC := gittest.WriteBlob(t, cfg, repoPath, []byte("c"))
+				dirATree := gittest.WriteTree(t, cfg, repoPath, []gittest.TreeEntry{
+					{
+						Mode: "100644",
+						Path: "file1InDirA",
+						OID:  blobA,
+					},
+				})
+
+				treeID := gittest.WriteTree(t, cfg, repoPath, []gittest.TreeEntry{
+					{
+						Mode: "100644",
+						Path: "fileB",
+						OID:  blobB,
+					},
+					{
+						Mode: "100644",
+						Path: "fileC",
+						OID:  blobC,
+					},
+					{
+						Mode: "040000",
+						Path: "dirA",
+						OID:  dirATree,
+					},
+				})
+
+				return gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("master"),
+						gittest.WithTree(treeID)),
+					TreeEntry{
+						OID:  treeID,
+						Mode: "040000",
+						Type: Tree,
+						Entries: []*TreeEntry{
+							{
+								Mode: "040000",
+								Path: "dirA",
+								Type: Tree,
+								OID:  dirATree,
+								Entries: []*TreeEntry{
+									{
+										Mode: "100644",
+										Path: "file1InDirA",
+										OID:  blobA,
+										Type: Blob,
+									},
+								},
+							},
+							{
+								Mode: "100644",
+								Path: "fileB",
+								OID:  blobB,
+								Type: Blob,
+							},
+							{
+								Mode: "100644",
+								Path: "fileC",
+								OID:  blobC,
+								Type: Blob,
+							},
+						},
+					}
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			repoProto, repoPath := gittest.CreateRepository(
+				t,
+				ctx,
+				cfg,
+				gittest.CreateRepositoryConfig{
+					SkipCreationViaService: true,
+				})
+			repo := NewTestRepo(t, cfg, repoProto)
+
+			commitID, expectedTree := tc.setupTree(t, repoPath)
+
+			fullTree, err := repo.GetFullTree(
+				ctx,
+				git.Revision(commitID),
+			)
+			require.NoError(t, err)
+			require.Equal(t, expectedTree, *fullTree)
+		})
+	}
+}