worktree: reset and checkout support for submodules

1 files changed, 42 insertions, 12 deletions

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 {