1 files changed, 46 insertions, 6 deletions

diff --git a/plumbing/object/tree.go b/plumbing/object/tree.go
index c2399f8..c36a137 100644
--- a/plumbing/object/tree.go
+++ b/plumbing/object/tree.go
@@ -2,10 +2,12 @@ package object
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"path"
+	"path/filepath"
 	"strings"
 
 	"gopkg.in/src-d/go-git.v4/plumbing"
@@ -24,6 +26,7 @@ var (
 	ErrMaxTreeDepth      = errors.New("maximum tree depth exceeded")
 	ErrFileNotFound      = errors.New("file not found")
 	ErrDirectoryNotFound = errors.New("directory not found")
+	ErrEntryNotFound     = errors.New("entry not found")
 )
 
 // Tree is basically like a directory - it references a bunch of other trees
@@ -34,6 +37,7 @@ type Tree struct {
 
 	s storer.EncodedObjectStorer
 	m map[string]*TreeEntry
+	t map[string]*Tree // tree path cache
 }
 
 // GetTree gets a tree from an object storer and decodes it.
@@ -111,14 +115,37 @@ func (t *Tree) TreeEntryFile(e *TreeEntry) (*File, error) {
 
 // FindEntry search a TreeEntry in this tree or any subtree.
 func (t *Tree) FindEntry(path string) (*TreeEntry, error) {
+	if t.t == nil {
+		t.t = make(map[string]*Tree)
+	}
+
 	pathParts := strings.Split(path, "/")
+	startingTree := t
+	pathCurrent := ""
+
+	// search for the longest path in the tree path cache
+	for i := len(pathParts); i > 1; i-- {
+		path := filepath.Join(pathParts[:i]...)
+
+		tree, ok := t.t[path]
+		if ok {
+			startingTree = tree
+			pathParts = pathParts[i:]
+			pathCurrent = path
+
+			break
+		}
+	}
 
 	var tree *Tree
 	var err error
-	for tree = t; len(pathParts) > 1; pathParts = pathParts[1:] {
+	for tree = startingTree; len(pathParts) > 1; pathParts = pathParts[1:] {
 		if tree, err = tree.dir(pathParts[0]); err != nil {
 			return nil, err
 		}
+
+		pathCurrent = filepath.Join(pathCurrent, pathParts[0])
+		t.t[pathCurrent] = tree
 	}
 
 	return tree.entry(pathParts[0])
@@ -141,8 +168,6 @@ func (t *Tree) dir(baseName string) (*Tree, error) {
 	return tree, err
 }
 
-var errEntryNotFound = errors.New("entry not found")
-
 func (t *Tree) entry(baseName string) (*TreeEntry, error) {
 	if t.m == nil {
 		t.buildMap()
@@ -150,7 +175,7 @@ func (t *Tree) entry(baseName string) (*TreeEntry, error) {
 
 	entry, ok := t.m[baseName]
 	if !ok {
-		return nil, errEntryNotFound
+		return nil, ErrEntryNotFound
 	}
 
 	return entry, nil
@@ -270,15 +295,30 @@ func (from *Tree) Diff(to *Tree) (Changes, error) {
 	return DiffTree(from, to)
 }
 
+// Diff returns a list of changes between this tree and the provided one
+// Error will be returned if context expires
+// Provided context must be non nil
+func (from *Tree) DiffContext(ctx context.Context, to *Tree) (Changes, error) {
+	return DiffTreeContext(ctx, from, to)
+}
+
 // Patch returns a slice of Patch objects with all the changes between trees
 // in chunks. This representation can be used to create several diff outputs.
 func (from *Tree) Patch(to *Tree) (*Patch, error) {
-	changes, err := DiffTree(from, to)
+	return from.PatchContext(context.Background(), to)
+}
+
+// Patch returns a slice of Patch objects with all the changes between trees
+// in chunks. This representation can be used to create several diff outputs.
+// If context expires, an error will be returned
+// Provided context must be non-nil
+func (from *Tree) PatchContext(ctx context.Context, to *Tree) (*Patch, error) {
+	changes, err := DiffTreeContext(ctx, from, to)
 	if err != nil {
 		return nil, err
 	}
 
-	return changes.Patch()
+	return changes.PatchContext(ctx)
 }
 
 // treeEntryIter facilitates iterating through the TreeEntry objects in a Tree.