worktree: reset and checkout support for submodules

2 files changed, 81 insertions, 12 deletions

diff --git a/utils/merkletrie/filesystem/node.go b/utils/merkletrie/filesystem/node.go
index 6c09d29..fc8f191 100644
--- a/utils/merkletrie/filesystem/node.go
+++ b/utils/merkletrie/filesystem/node.go
@@ -21,7 +21,8 @@ var ignore = map[string]bool{
 // This implementation implements a "standard" hash method being able to be
 // compared with any other noder.Noder implementation inside of go-git.
 type node struct {
-	fs billy.Filesystem
+	fs         billy.Filesystem
+	submodules map[string]plumbing.Hash
 
 	path     string
 	hash     []byte
@@ -29,9 +30,16 @@ type node struct {
 	isDir    bool
 }
 
-// NewRootNode returns the root node based on a given billy.Filesystem
-func NewRootNode(fs billy.Filesystem) noder.Noder {
-	return &node{fs: fs, isDir: true}
+// NewRootNode returns the root node based on a given billy.Filesystem.
+//
+// In order to provide the submodule hash status, a map[string]plumbing.Hash
+// should be provided where the key is the path of the submodule and the commit
+// of the submodule HEAD
+func NewRootNode(
+	fs billy.Filesystem,
+	submodules map[string]plumbing.Hash,
+) noder.Noder {
+	return &node{fs: fs, submodules: submodules, isDir: true}
 }
 
 // Hash the hash of a filesystem is the result of concatenating the computed
@@ -100,17 +108,27 @@ func (n *node) calculateChildren() error {
 
 func (n *node) newChildNode(file billy.FileInfo) (*node, error) {
 	path := filepath.Join(n.path, file.Name())
+
 	hash, err := n.calculateHash(path, file)
 	if err != nil {
 		return nil, err
 	}
 
-	return &node{
-		fs:    n.fs,
+	node := &node{
+		fs:         n.fs,
+		submodules: n.submodules,
+
 		path:  path,
 		hash:  hash,
 		isDir: file.IsDir(),
-	}, nil
+	}
+
+	if hash, isSubmodule := n.submodules[path]; isSubmodule {
+		node.hash = append(hash[:], filemode.Submodule.Bytes()...)
+		node.isDir = false
+	}
+
+	return node, nil
 }
 
 func (n *node) calculateHash(path string, file billy.FileInfo) ([]byte, error) {
diff --git a/utils/merkletrie/filesystem/node_test.go b/utils/merkletrie/filesystem/node_test.go
index b7c124d..f783c15 100644
--- a/utils/merkletrie/filesystem/node_test.go
+++ b/utils/merkletrie/filesystem/node_test.go
@@ -9,6 +9,7 @@ import (
 	. "gopkg.in/check.v1"
 	"gopkg.in/src-d/go-billy.v2"
 	"gopkg.in/src-d/go-billy.v2/memfs"
+	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie"
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie/noder"
 )
@@ -30,7 +31,12 @@ func (s *NoderSuite) TestDiff(c *C) {
 	WriteFile(fsB, "qux/bar", []byte("foo"), 0644)
 	WriteFile(fsB, "qux/qux", []byte("foo"), 0644)
 
-	ch, err := merkletrie.DiffTree(NewRootNode(fsA), NewRootNode(fsB), IsEquals)
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, nil),
+		NewRootNode(fsB, nil),
+		IsEquals,
+	)
+
 	c.Assert(err, IsNil)
 	c.Assert(ch, HasLen, 0)
 }
@@ -46,7 +52,12 @@ func (s *NoderSuite) TestDiffChangeContent(c *C) {
 	WriteFile(fsB, "qux/bar", []byte("bar"), 0644)
 	WriteFile(fsB, "qux/qux", []byte("foo"), 0644)
 
-	ch, err := merkletrie.DiffTree(NewRootNode(fsA), NewRootNode(fsB), IsEquals)
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, nil),
+		NewRootNode(fsB, nil),
+		IsEquals,
+	)
+
 	c.Assert(err, IsNil)
 	c.Assert(ch, HasLen, 1)
 }
@@ -58,7 +69,12 @@ func (s *NoderSuite) TestDiffChangeMissing(c *C) {
 	fsB := memfs.New()
 	WriteFile(fsB, "bar", []byte("bar"), 0644)
 
-	ch, err := merkletrie.DiffTree(NewRootNode(fsA), NewRootNode(fsB), IsEquals)
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, nil),
+		NewRootNode(fsB, nil),
+		IsEquals,
+	)
+
 	c.Assert(err, IsNil)
 	c.Assert(ch, HasLen, 2)
 }
@@ -70,7 +86,12 @@ func (s *NoderSuite) TestDiffChangeMode(c *C) {
 	fsB := memfs.New()
 	WriteFile(fsB, "foo", []byte("foo"), 0755)
 
-	ch, err := merkletrie.DiffTree(NewRootNode(fsA), NewRootNode(fsB), IsEquals)
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, nil),
+		NewRootNode(fsB, nil),
+		IsEquals,
+	)
+
 	c.Assert(err, IsNil)
 	c.Assert(ch, HasLen, 1)
 }
@@ -82,11 +103,41 @@ func (s *NoderSuite) TestDiffChangeModeNotRelevant(c *C) {
 	fsB := memfs.New()
 	WriteFile(fsB, "foo", []byte("foo"), 0655)
 
-	ch, err := merkletrie.DiffTree(NewRootNode(fsA), NewRootNode(fsB), IsEquals)
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, nil),
+		NewRootNode(fsB, nil),
+		IsEquals,
+	)
+
 	c.Assert(err, IsNil)
 	c.Assert(ch, HasLen, 0)
 }
 
+func (s *NoderSuite) TestDiffDirectory(c *C) {
+	fsA := memfs.New()
+	fsA.MkdirAll("qux/bar", 0644)
+
+	fsB := memfs.New()
+	fsB.MkdirAll("qux/bar", 0644)
+
+	ch, err := merkletrie.DiffTree(
+		NewRootNode(fsA, map[string]plumbing.Hash{
+			"qux/bar": plumbing.NewHash("aa102815663d23f8b75a47e7a01965dcdc96468c"),
+		}),
+		NewRootNode(fsB, map[string]plumbing.Hash{
+			"qux/bar": plumbing.NewHash("19102815663d23f8b75a47e7a01965dcdc96468c"),
+		}),
+		IsEquals,
+	)
+
+	c.Assert(err, IsNil)
+	c.Assert(ch, HasLen, 1)
+
+	a, err := ch[0].Action()
+	c.Assert(err, IsNil)
+	c.Assert(a, Equals, merkletrie.Modify)
+}
+
 func WriteFile(fs billy.Filesystem, filename string, data []byte, perm os.FileMode) error {
 	f, err := fs.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
 	if err != nil {