2 files changed, 82 insertions, 0 deletions
diff --git a/utils/merkletrie/difftree.go b/utils/merkletrie/difftree.go
index 2294096..d57ed13 100644
--- a/utils/merkletrie/difftree.go
+++ b/utils/merkletrie/difftree.go
@@ -248,15 +248,30 @@ package merkletrie
 // h: else of i
 
 import (
+	"context"
+	"errors"
 	"fmt"
 
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie/noder"
 )
 
+var (
+	ErrCanceled = errors.New("operation canceled")
+)
+
 // DiffTree calculates the list of changes between two merkletries.  It
 // uses the provided hashEqual callback to compare noders.
 func DiffTree(fromTree, toTree noder.Noder,
 	hashEqual noder.Equal) (Changes, error) {
+	return DiffTreeContext(context.Background(), fromTree, toTree, hashEqual)
+}
+
+// DiffTree calculates the list of changes between two merkletries.  It
+// uses the provided hashEqual callback to compare noders.
+// Error will be returned if context expires
+// Provided context must be non nil
+func DiffTreeContext(ctx context.Context, fromTree, toTree noder.Noder,
+	hashEqual noder.Equal) (Changes, error) {
 	ret := NewChanges()
 
 	ii, err := newDoubleIter(fromTree, toTree, hashEqual)
@@ -265,6 +280,12 @@ func DiffTree(fromTree, toTree noder.Noder,
 	}
 
 	for {
+		select {
+		case <-ctx.Done():
+			return nil, ErrCanceled
+		default:
+		}
+
 		from := ii.from.current
 		to := ii.to.current
 
diff --git a/utils/merkletrie/difftree_test.go b/utils/merkletrie/difftree_test.go
index 9f033b1..ab0eb57 100644
--- a/utils/merkletrie/difftree_test.go
+++ b/utils/merkletrie/difftree_test.go
@@ -2,6 +2,7 @@ package merkletrie_test
 
 import (
 	"bytes"
+	ctx "context"
 	"fmt"
 	"reflect"
 	"sort"
@@ -61,9 +62,45 @@ func (t diffTreeTest) innerRun(c *C, context string, reverse bool) {
 	c.Assert(obtained, changesEquals, expected, comment)
 }
 
+func (t diffTreeTest) innerRunCtx(c *C, context string, reverse bool) {
+	comment := Commentf("\n%s", context)
+	if reverse {
+		comment = Commentf("%s [REVERSED]", comment.CheckCommentString())
+	}
+
+	a, err := fsnoder.New(t.from)
+	c.Assert(err, IsNil, comment)
+	comment = Commentf("%s\n\t    from = %s", comment.CheckCommentString(), a)
+
+	b, err := fsnoder.New(t.to)
+	c.Assert(err, IsNil, comment)
+	comment = Commentf("%s\n\t      to = %s", comment.CheckCommentString(), b)
+
+	expected, err := newChangesFromString(t.expected)
+	c.Assert(err, IsNil, comment)
+
+	if reverse {
+		a, b = b, a
+		expected = expected.reverse()
+	}
+	comment = Commentf("%s\n\texpected = %s", comment.CheckCommentString(), expected)
+
+	results, err := merkletrie.DiffTreeContext(ctx.Background(), a, b, fsnoder.HashEqual)
+	c.Assert(err, IsNil, comment)
+
+	obtained, err := newChanges(results)
+	c.Assert(err, IsNil, comment)
+
+	comment = Commentf("%s\n\tobtained = %s", comment.CheckCommentString(), obtained)
+
+	c.Assert(obtained, changesEquals, expected, comment)
+}
+
 func (t diffTreeTest) run(c *C, context string) {
 	t.innerRun(c, context, false)
 	t.innerRun(c, context, true)
+	t.innerRunCtx(c, context, false)
+	t.innerRunCtx(c, context, true)
 }
 
 type change struct {
@@ -437,3 +474,27 @@ func (s *DiffTreeSuite) TestIssue275(c *C) {
 		},
 	})
 }
+
+func (s *DiffTreeSuite) TestCancel(c *C) {
+	t :=  diffTreeTest{"()", "(a<> b<1> c() d<> e<2> f())", "+a +b +d +e"}
+	comment := Commentf("\n%s", "test cancel:")
+
+	a, err := fsnoder.New(t.from)
+	c.Assert(err, IsNil, comment)
+	comment = Commentf("%s\n\t    from = %s", comment.CheckCommentString(), a)
+
+	b, err := fsnoder.New(t.to)
+	c.Assert(err, IsNil, comment)
+	comment = Commentf("%s\n\t      to = %s", comment.CheckCommentString(), b)
+
+	expected, err := newChangesFromString(t.expected)
+	c.Assert(err, IsNil, comment)
+
+	comment = Commentf("%s\n\texpected = %s", comment.CheckCommentString(), expected)
+	context, cancel := ctx.WithCancel(ctx.Background())
+	cancel()
+	results, err := merkletrie.DiffTreeContext(context, a, b, fsnoder.HashEqual)
+	c.Assert(results, IsNil, comment)
+	c.Assert(err, ErrorMatches, "operation canceled")
+
+}