worktree: reset and checkout support for submodules

8 files changed, 444 insertions, 53 deletions

diff --git a/storage/storer.go b/storage/storer.go
index 50bfd55..d1a94f2 100644
--- a/storage/storer.go
+++ b/storage/storer.go
@@ -20,7 +20,7 @@ type Storer interface {
 
 // ModuleStorer allows interact with the modules' Storers
 type ModuleStorer interface {
-	// Module returns a Storer reprensting a submodule, if not exists returns a
+	// Module returns a Storer representing a submodule, if not exists returns a
 	// new empty Storer is returned
 	Module(name string) (Storer, error)
 }
diff --git a/submodule.go b/submodule.go
index c711a2b..fd23b9d 100644
--- a/submodule.go
+++ b/submodule.go
@@ -1,10 +1,13 @@
 package git
 
 import (
+	"bytes"
 	"errors"
+	"fmt"
 
 	"gopkg.in/src-d/go-git.v4/config"
 	"gopkg.in/src-d/go-git.v4/plumbing"
+	"gopkg.in/src-d/go-git.v4/plumbing/format/index"
 )
 
 var (
@@ -15,6 +18,7 @@ var (
 // Submodule a submodule allows you to keep another Git repository in a
 // subdirectory of your repository.
 type Submodule struct {
+	// initialized defines if a submodule was already initialized.
 	initialized bool
 
 	c *config.Submodule
@@ -26,7 +30,7 @@ func (s *Submodule) Config() *config.Submodule {
 	return s.c
 }
 
-// Init initialize the submodule reading the recoreded Entry in the index for
+// Init initialize the submodule reading the recorded Entry in the index for
 // the given submodule
 func (s *Submodule) Init() error {
 	cfg, err := s.w.r.Storer.Config()
@@ -45,8 +49,54 @@ func (s *Submodule) Init() error {
 	return s.w.r.Storer.SetConfig(cfg)
 }
 
+// Status returns the status of the submodule.
+func (s *Submodule) Status() (*SubmoduleStatus, error) {
+	idx, err := s.w.r.Storer.Index()
+	if err != nil {
+		return nil, err
+	}
+
+	return s.status(idx)
+}
+
+func (s *Submodule) status(idx *index.Index) (*SubmoduleStatus, error) {
+	e, err := idx.Entry(s.c.Path)
+	if err != nil {
+		return nil, err
+	}
+
+	status := &SubmoduleStatus{
+		Path:     s.c.Path,
+		Expected: e.Hash,
+	}
+
+	if !s.initialized {
+		return status, nil
+	}
+
+	r, err := s.Repository()
+	if err != nil {
+		return nil, err
+	}
+
+	head, err := r.Head()
+	if err == nil {
+		status.Current = head.Hash()
+	}
+
+	if err != nil && err == plumbing.ErrReferenceNotFound {
+		err = nil
+	}
+
+	return status, err
+}
+
 // Repository returns the Repository represented by this submodule
 func (s *Submodule) Repository() (*Repository, error) {
+	if !s.initialized {
+		return nil, ErrSubmoduleNotInitialized
+	}
+
 	storer, err := s.w.r.Storer.Module(s.c.Name)
 	if err != nil {
 		return nil, err
@@ -76,9 +126,13 @@ func (s *Submodule) Repository() (*Repository, error) {
 }
 
 // Update the registered submodule to match what the superproject expects, the
-// submodule should be initilized first calling the Init method or setting in
+// submodule should be initialized first calling the Init method or setting in
 // the options SubmoduleUpdateOptions.Init equals true
 func (s *Submodule) Update(o *SubmoduleUpdateOptions) error {
+	return s.update(o, plumbing.ZeroHash)
+}
+
+func (s *Submodule) update(o *SubmoduleUpdateOptions, forceHash plumbing.Hash) error {
 	if !s.initialized && !o.Init {
 		return ErrSubmoduleNotInitialized
 	}
@@ -89,17 +143,27 @@ func (s *Submodule) Update(o *SubmoduleUpdateOptions) error {
 		}
 	}
 
-	e, err := s.w.readIndexEntry(s.c.Path)
+	idx, err := s.w.r.Storer.Index()
 	if err != nil {
 		return err
 	}
 
+	hash := forceHash
+	if hash.IsZero() {
+		e, err := idx.Entry(s.c.Path)
+		if err != nil {
+			return err
+		}
+
+		hash = e.Hash
+	}
+
 	r, err := s.Repository()
 	if err != nil {
 		return err
 	}
 
-	if err := s.fetchAndCheckout(r, o, e.Hash); err != nil {
+	if err := s.fetchAndCheckout(r, o, hash); err != nil {
 		return err
 	}
 
@@ -123,6 +187,7 @@ func (s *Submodule) doRecursiveUpdate(r *Repository, o *SubmoduleUpdateOptions)
 
 	new := &SubmoduleUpdateOptions{}
 	*new = *o
+
 	new.RecurseSubmodules--
 	return l.Update(new)
 }
@@ -148,10 +213,10 @@ func (s *Submodule) fetchAndCheckout(r *Repository, o *SubmoduleUpdateOptions, h
 	return r.Storer.SetReference(head)
 }
 
-// Submodules list of several submodules from the same repository
+// Submodules list of several submodules from the same repository.
 type Submodules []*Submodule
 
-// Init initializes the submodules in this list
+// Init initializes the submodules in this list.
 func (s Submodules) Init() error {
 	for _, sub := range s {
 		if err := sub.Init(); err != nil {
@@ -162,7 +227,7 @@ func (s Submodules) Init() error {
 	return nil
 }
 
-// Update updates all the submodules in this list
+// Update updates all the submodules in this list.
 func (s Submodules) Update(o *SubmoduleUpdateOptions) error {
 	for _, sub := range s {
 		if err := sub.Update(o); err != nil {
@@ -172,3 +237,85 @@ func (s Submodules) Update(o *SubmoduleUpdateOptions) error {
 
 	return nil
 }
+
+// Status returns the status of the submodules.
+func (s Submodules) Status() (SubmodulesStatus, error) {
+	var list SubmodulesStatus
+
+	var r *Repository
+	for _, sub := range s {
+		if r == nil {
+			r = sub.w.r
+		}
+
+		idx, err := r.Storer.Index()
+		if err != nil {
+			return nil, err
+		}
+
+		status, err := sub.status(idx)
+		if err != nil {
+			return nil, err
+		}
+
+		list = append(list, status)
+	}
+
+	return list, nil
+}
+
+// SubmodulesStatus contains the status for all submodiles in the worktree
+type SubmodulesStatus []*SubmoduleStatus
+
+// String is equivalent to `git submodule status`
+func (s SubmodulesStatus) String() string {
+	buf := bytes.NewBuffer(nil)
+	for _, sub := range s {
+		fmt.Fprintln(buf, sub)
+	}
+
+	return buf.String()
+}
+
+// SubmoduleStatus contains the status for a submodule in the worktree
+type SubmoduleStatus struct {
+	Path     string
+	Current  plumbing.Hash
+	Expected plumbing.Hash
+	Branch   plumbing.ReferenceName
+}
+
+// IsClean is the HEAD of the submodule is equals to the expected commit
+func (s *SubmoduleStatus) IsClean() bool {
+	return s.Current == s.Expected
+}
+
+// String is equivalent to `git submodule status <submodule>`
+//
+// This will print the SHA-1 of the currently checked out commit for a
+// submodule, along with the submodule path and the output of git describe fo
+// the SHA-1. Each SHA-1 will be prefixed with - if the submodule is not
+// initialized, + if the currently checked out submodule commit does not match
+// the SHA-1 found in the index of the containing repository.
+func (s *SubmoduleStatus) String() string {
+	var extra string
+	var status = ' '
+
+	if s.Current.IsZero() {
+		status = '-'
+	} else if !s.IsClean() {
+		status = '+'
+	}
+
+	if len(s.Branch) != 0 {
+		extra = string(s.Branch[5:])
+	} else if !s.Current.IsZero() {
+		extra = s.Current.String()[:7]
+	}
+
+	if extra != "" {
+		extra = fmt.Sprintf(" (%s)", extra)
+	}
+
+	return fmt.Sprintf("%c%s %s%s", status, s.Expected, s.Path, extra)
+}
diff --git a/submodule_test.go b/submodule_test.go
index 88afa18..6e06191 100644
--- a/submodule_test.go
+++ b/submodule_test.go
@@ -48,14 +48,21 @@ func (s *SubmoduleSuite) TestInit(c *C) {
 	sm, err := s.Worktree.Submodule("basic")
 	c.Assert(err, IsNil)
 
+	c.Assert(sm.initialized, Equals, false)
 	err = sm.Init()
 	c.Assert(err, IsNil)
 
+	c.Assert(sm.initialized, Equals, true)
+
 	cfg, err := s.Repository.Config()
 	c.Assert(err, IsNil)
 
 	c.Assert(cfg.Submodules, HasLen, 1)
 	c.Assert(cfg.Submodules["basic"], NotNil)
+
+	status, err := sm.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status.IsClean(), Equals, false)
 }
 
 func (s *SubmoduleSuite) TestUpdate(c *C) {
@@ -74,6 +81,19 @@ func (s *SubmoduleSuite) TestUpdate(c *C) {
 	ref, err := r.Reference(plumbing.HEAD, true)
 	c.Assert(err, IsNil)
 	c.Assert(ref.Hash().String(), Equals, "6ecf0ef2c2dffb796033e5a02219af86ec6584e5")
+
+	status, err := sm.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status.IsClean(), Equals, true)
+}
+
+func (s *SubmoduleSuite) TestRepositoryWithoutInit(c *C) {
+	sm, err := s.Worktree.Submodule("basic")
+	c.Assert(err, IsNil)
+
+	r, err := sm.Repository()
+	c.Assert(err, Equals, ErrSubmoduleNotInitialized)
+	c.Assert(r, IsNil)
 }
 
 func (s *SubmoduleSuite) TestUpdateWithoutInit(c *C) {
@@ -161,3 +181,12 @@ func (s *SubmoduleSuite) TestSubmodulesInit(c *C) {
 		c.Assert(m.initialized, Equals, true)
 	}
 }
+
+func (s *SubmoduleSuite) TestSubmodulesStatus(c *C) {
+	sm, err := s.Worktree.Submodules()
+	c.Assert(err, IsNil)
+
+	status, err := sm.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 2)
+}
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 {
diff --git a/worktree.go b/worktree.go
index ec8fad2..e92449c 100644
--- a/worktree.go
+++ b/worktree.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 
 	"gopkg.in/src-d/go-git.v4/config"
 	"gopkg.in/src-d/go-git.v4/plumbing"
@@ -207,31 +208,103 @@ func (w *Worktree) checkoutChange(ch merkletrie.Change, t *object.Tree, idx *ind
 		return err
 	}
 
+	var e *object.TreeEntry
+	var name string
+	var isSubmodule bool
+
+	switch a {
+	case merkletrie.Modify, merkletrie.Insert:
+		name = ch.To.String()
+		e, err = t.FindEntry(name)
+		if err != nil {
+			return err
+		}
+
+		isSubmodule = e.Mode == filemode.Submodule
+	case merkletrie.Delete:
+		name = ch.From.String()
+		ie, err := idx.Entry(name)
+		if err != nil {
+			return err
+		}
+
+		isSubmodule = ie.Mode == filemode.Submodule
+	}
+
+	if isSubmodule {
+		return w.checkoutChangeSubmodule(name, a, e, idx)
+	}
+
+	return w.checkoutChangeRegularFile(name, a, t, e, idx)
+}
+
+func (w *Worktree) checkoutChangeSubmodule(name string,
+	a merkletrie.Action,
+	e *object.TreeEntry,
+	idx *index.Index,
+) error {
 	switch a {
 	case merkletrie.Modify:
-		name := ch.To.String()
+		sub, err := w.Submodule(name)
+		if err != nil {
+			return err
+		}
+
+		if !sub.initialized {
+			return nil
+		}
+
 		if err := w.rmIndexFromFile(name, idx); err != nil {
 			return err
 		}
 
-		// to apply perm changes the file is deleted, billy doesn't implement
-		// chmod
-		if err := w.fs.Remove(name); err != nil {
+		if err := w.addIndexFromTreeEntry(name, e, idx); err != nil {
 			return err
 		}
 
-		fallthrough
+		return sub.update(&SubmoduleUpdateOptions{}, e.Hash)
 	case merkletrie.Insert:
-		name := ch.To.String()
-		e, err := t.FindEntry(name)
+		mode, err := e.Mode.ToOSFileMode()
 		if err != nil {
 			return err
 		}
 
-		if e.Mode == filemode.Submodule {
-			return w.addIndexFromTreeEntry(name, e, idx)
+		if err := w.fs.MkdirAll(name, mode); err != nil {
+			return err
 		}
 
+		return w.addIndexFromTreeEntry(name, e, idx)
+	case merkletrie.Delete:
+		if err := rmFileAndDirIfEmpty(w.fs, name); err != nil {
+			return err
+		}
+
+		return w.rmIndexFromFile(name, idx)
+	}
+
+	return nil
+}
+
+func (w *Worktree) checkoutChangeRegularFile(name string,
+	a merkletrie.Action,
+	t *object.Tree,
+	e *object.TreeEntry,
+	idx *index.Index,
+) error {
+	switch a {
+	case merkletrie.Modify:
+		if err := w.rmIndexFromFile(name, idx); err != nil {
+			return err
+		}
+
+		// to apply perm changes the file is deleted, billy doesn't implement
+		// chmod
+		if err := w.fs.Remove(name); err != nil {
+			return err
+		}
+
+		fallthrough
+	case merkletrie.Insert:
 		f, err := t.File(name)
 		if err != nil {
 			return err
@@ -243,8 +316,7 @@ func (w *Worktree) checkoutChange(ch merkletrie.Change, t *object.Tree, idx *ind
 
 		return w.addIndexFromFile(name, e.Hash, idx)
 	case merkletrie.Delete:
-		name := ch.From.String()
-		if err := w.fs.Remove(name); err != nil {
+		if err := rmFileAndDirIfEmpty(w.fs, name); err != nil {
 			return err
 		}
 
@@ -413,19 +485,20 @@ func (w *Worktree) readGitmodulesFile() (*config.Modules, error) {
 	return m, m.Unmarshal(input)
 }
 
-func (w *Worktree) readIndexEntry(path string) (index.Entry, error) {
-	var e index.Entry
+func rmFileAndDirIfEmpty(fs billy.Filesystem, name string) error {
+	if err := billy.RemoveAll(fs, name); err != nil {
+		return err
+	}
 
-	idx, err := w.r.Storer.Index()
+	path := filepath.Dir(name)
+	files, err := fs.ReadDir(path)
 	if err != nil {
-		return e, err
+		return err
 	}
 
-	for _, e := range idx.Entries {
-		if e.Name == path {
-			return e, nil
-		}
+	if len(files) == 0 {
+		fs.Remove(path)
 	}
 
-	return e, fmt.Errorf("unable to find %q entry in the index", path)
+	return nil
 }
diff --git a/worktree_status.go b/worktree_status.go
index ae7518e..373f161 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -28,12 +28,12 @@ func (w *Worktree) Status() (Status, error) {
 func (w *Worktree) status(commit plumbing.Hash) (Status, error) {
 	s := make(Status, 0)
 
-	right, err := w.diffStagingWithWorktree()
+	left, err := w.diffCommitWithStaging(commit, false)
 	if err != nil {
 		return nil, err
 	}
 
-	for _, ch := range right {
+	for _, ch := range left {
 		a, err := ch.Action()
 		if err != nil {
 			return nil, err
@@ -41,21 +41,20 @@ func (w *Worktree) status(commit plumbing.Hash) (Status, error) {
 
 		switch a {
 		case merkletrie.Delete:
-			s.File(ch.From.String()).Worktree = Deleted
+			s.File(ch.From.String()).Staging = Deleted
 		case merkletrie.Insert:
-			s.File(ch.To.String()).Worktree = Untracked
-			s.File(ch.To.String()).Staging = Untracked
+			s.File(ch.To.String()).Staging = Added
 		case merkletrie.Modify:
-			s.File(ch.To.String()).Worktree = Modified
+			s.File(ch.To.String()).Staging = Modified
 		}
 	}
 
-	left, err := w.diffCommitWithStaging(commit, false)
+	right, err := w.diffStagingWithWorktree()
 	if err != nil {
 		return nil, err
 	}
 
-	for _, ch := range left {
+	for _, ch := range right {
 		a, err := ch.Action()
 		if err != nil {
 			return nil, err
@@ -63,11 +62,12 @@ func (w *Worktree) status(commit plumbing.Hash) (Status, error) {
 
 		switch a {
 		case merkletrie.Delete:
-			s.File(ch.From.String()).Staging = Deleted
+			s.File(ch.From.String()).Worktree = Deleted
 		case merkletrie.Insert:
-			s.File(ch.To.String()).Staging = Added
+			s.File(ch.To.String()).Worktree = Untracked
+			s.File(ch.To.String()).Staging = Untracked
 		case merkletrie.Modify:
-			s.File(ch.To.String()).Staging = Modified
+			s.File(ch.To.String()).Worktree = Modified
 		}
 	}
 
@@ -81,10 +81,40 @@ func (w *Worktree) diffStagingWithWorktree() (merkletrie.Changes, error) {
 	}
 
 	from := index.NewRootNode(idx)
-	to := filesystem.NewRootNode(w.fs)
+	submodules, err := w.getSubmodulesStatus()
+	if err != nil {
+		return nil, err
+	}
+
+	to := filesystem.NewRootNode(w.fs, submodules)
 	return merkletrie.DiffTree(from, to, diffTreeIsEquals)
 }
 
+func (w *Worktree) getSubmodulesStatus() (map[string]plumbing.Hash, error) {
+	o := map[string]plumbing.Hash{}
+
+	sub, err := w.Submodules()
+	if err != nil {
+		return nil, err
+	}
+
+	status, err := sub.Status()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, s := range status {
+		if s.Current.IsZero() {
+			o[s.Path] = s.Expected
+			continue
+		}
+
+		o[s.Path] = s.Current
+	}
+
+	return o, nil
+}
+
 func (w *Worktree) diffCommitWithStaging(commit plumbing.Hash, reverse bool) (merkletrie.Changes, error) {
 	idx, err := w.r.Storer.Index()
 	if err != nil {
diff --git a/worktree_test.go b/worktree_test.go
index 59197a6..3393469 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -54,6 +54,48 @@ func (s *WorktreeSuite) TestCheckout(c *C) {
 	c.Assert(idx.Entries, HasLen, 9)
 }
 
+func (s *WorktreeSuite) TestCheckoutSubmodule(c *C) {
+	url := "https://github.com/git-fixtures/submodule.git"
+	w := &Worktree{
+		r:  s.NewRepository(fixtures.ByURL(url).One()),
+		fs: memfs.New(),
+	}
+
+	// we delete the index, since the fixture comes with a real index
+	err := w.r.Storer.SetIndex(&index.Index{Version: 2})
+	c.Assert(err, IsNil)
+
+	err = w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status.IsClean(), Equals, true)
+}
+
+func (s *WorktreeSuite) TestCheckoutSubmoduleInitialized(c *C) {
+	url := "https://github.com/git-fixtures/submodule.git"
+	w := &Worktree{
+		r:  s.NewRepository(fixtures.ByURL(url).One()),
+		fs: memfs.New(),
+	}
+
+	err := w.r.Storer.SetIndex(&index.Index{Version: 2})
+	c.Assert(err, IsNil)
+
+	err = w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+	sub, err := w.Submodules()
+	c.Assert(err, IsNil)
+
+	err = sub.Update(&SubmoduleUpdateOptions{Init: true})
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status.IsClean(), Equals, true)
+}
+
 func (s *WorktreeSuite) TestCheckoutIndexMem(c *C) {
 	fs := memfs.New()
 	w := &Worktree{
@@ -119,6 +161,7 @@ func (s *WorktreeSuite) TestCheckoutChange(c *C) {
 
 	err := w.Checkout(&CheckoutOptions{})
 	c.Assert(err, IsNil)
+
 	head, err := w.r.Head()
 	c.Assert(err, IsNil)
 	c.Assert(head.Name().String(), Equals, "refs/heads/master")
@@ -143,6 +186,7 @@ func (s *WorktreeSuite) TestCheckoutChange(c *C) {
 
 	_, err = fs.Stat("README")
 	c.Assert(err, Equals, nil)
+
 	_, err = fs.Stat("vendor")
 	c.Assert(err, Equals, os.ErrNotExist)
 
@@ -200,7 +244,6 @@ func (s *WorktreeSuite) TestCheckoutBisect(c *C) {
 }
 
 func (s *WorktreeSuite) TestCheckoutBisectSubmodules(c *C) {
-	c.Skip("not-submodule-support")
 	s.testCheckoutBisect(c, "https://github.com/git-fixtures/submodule.git")
 }