worktree, reset implementation and status improvements

8 files changed, 265 insertions, 122 deletions

diff --git a/options.go b/options.go
index f46bb6d..50fba10 100644
--- a/options.go
+++ b/options.go
@@ -178,12 +178,15 @@ type SubmoduleUpdateOptions struct {
 	RecurseSubmodules SubmoduleRescursivity
 }
 
-// CheckoutOptions describes how a checkout operation should be performed.
+// CheckoutOptions describes how a checkout 31operation should be performed.
 type CheckoutOptions struct {
-	// Branch to be checked out, if empty uses `master`
+	// Hash to be checked out, if used HEAD will in detached mode. Branch and
+	// Hash are mutual exclusive.
+	Hash plumbing.Hash
+	// Branch to be checked out, if Branch and Hash are empty is set to `master`.
 	Branch plumbing.ReferenceName
-	Hash   plumbing.Hash
-	// RemoteName is the name of the remote to be pushed to.
+	// Force, if true when switching branches, proceed even if the index or the
+	// working tree differs from HEAD. This is used to throw away local changes
 	Force bool
 }
 
@@ -196,23 +199,34 @@ func (o *CheckoutOptions) Validate() error {
 	return nil
 }
 
-type ResetMode int
+// ResetMode defines the mode of a reset operation.
+type ResetMode int8
 
 const (
 	// HardReset resets the index and working tree. Any changes to tracked files
 	// in the working tree are discarded.
 	HardReset ResetMode = iota
-	// MixedReset Resets the index but not the working tree (i.e., the changed
+	// MixedReset resets the index but not the working tree (i.e., the changed
 	// files are preserved but not marked for commit) and reports what has not
 	// been updated. This is the default action.
 	MixedReset
+	// MergeReset resets the index and updates the files in the working tree
+	// that are different between Commit and HEAD, but keeps those which are
+	// different between the index and working tree (i.e. which have changes
+	// which have not been added).
+	//
+	// If a file that is different between Commit and the index has unstaged
+	// changes, reset is aborted.
+	MergeReset
 )
 
 // ResetOptions describes how a reset operation should be performed.
 type ResetOptions struct {
 	// Commit, if commit is pressent set the current branch head (HEAD) to it.
 	Commit plumbing.Hash
-	// Mode
+	// Mode, form resets the current branch head to Commit and possibly updates
+	// the index (resetting it to the tree of Commit) and the working tree
+	// depending on Mode. If empty MixedReset is used.
 	Mode ResetMode
 }
 
diff --git a/plumbing/format/index/index.go b/plumbing/format/index/index.go
index 3675c4e..61e7d66 100644
--- a/plumbing/format/index/index.go
+++ b/plumbing/format/index/index.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"time"
 
+	"bytes"
+
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/filemode"
 )
@@ -50,12 +52,12 @@ type Index struct {
 
 // String is equivalent to `git ls-files --stage --debug`
 func (i *Index) String() string {
-	var o string
+	buf := bytes.NewBuffer(nil)
 	for _, e := range i.Entries {
-		o += e.String()
+		buf.WriteString(e.String())
 	}
 
-	return o
+	return buf.String()
 }
 
 // Entry represents a single file (or stage of a file) in the cache. An entry
@@ -90,15 +92,16 @@ type Entry struct {
 }
 
 func (e Entry) String() string {
-	var o string
-	o += fmt.Sprintf("%06o %s %d\t%s\n", e.Mode, e.Hash, e.Stage, e.Name)
-	o += fmt.Sprintf("  ctime: %d:%d\n", e.CreatedAt.Unix(), e.CreatedAt.Nanosecond())
-	o += fmt.Sprintf("  mtime: %d:%d\n", e.ModifiedAt.Unix(), e.ModifiedAt.Nanosecond())
-	o += fmt.Sprintf("  dev: %d\tino: %d\n", e.Dev, e.Inode)
-	o += fmt.Sprintf("  uid: %d\tgid: %d\n", e.UID, e.GID)
-	o += fmt.Sprintf("  size: %d\tflags: %x\n", e.Size, 0)
-
-	return o
+	buf := bytes.NewBuffer(nil)
+
+	fmt.Fprintf(buf, "%06o %s %d\t%s\n", e.Mode, e.Hash, e.Stage, e.Name)
+	fmt.Fprintf(buf, "  ctime: %d:%d\n", e.CreatedAt.Unix(), e.CreatedAt.Nanosecond())
+	fmt.Fprintf(buf, "  mtime: %d:%d\n", e.ModifiedAt.Unix(), e.ModifiedAt.Nanosecond())
+	fmt.Fprintf(buf, "  dev: %d\tino: %d\n", e.Dev, e.Inode)
+	fmt.Fprintf(buf, "  uid: %d\tgid: %d\n", e.UID, e.GID)
+	fmt.Fprintf(buf, "  size: %d\tflags: %x\n", e.Size, 0)
+
+	return buf.String()
 }
 
 // Tree contains pre-computed hashes for trees that can be derived from the
diff --git a/plumbing/object/tree.go b/plumbing/object/tree.go
index b768b96..d2265a8 100644
--- a/plumbing/object/tree.go
+++ b/plumbing/object/tree.go
@@ -109,7 +109,7 @@ func (t *Tree) TreeEntryFile(e *TreeEntry) (*File, error) {
 	return NewFile(e.Name, e.Mode, blob), nil
 }
 
-// FindEntry search a TreeEntry in this tree or any subtree
+// FindEntry search a TreeEntry in this tree or any subtree.
 func (t *Tree) FindEntry(path string) (*TreeEntry, error) {
 	pathParts := strings.Split(path, "/")
 
diff --git a/repository.go b/repository.go
index 2fdbafc..bb59afe 100644
--- a/repository.go
+++ b/repository.go
@@ -575,7 +575,7 @@ func (r *Repository) updateWorktree(branch plumbing.ReferenceName) error {
 		return err
 	}
 
-	return w.reset(&ResetOptions{
+	return w.Reset(&ResetOptions{
 		Commit: b.Hash(),
 	})
 }
diff --git a/status.go b/status.go
index e789f4a..2517e50 100644
--- a/status.go
+++ b/status.go
@@ -1,19 +1,23 @@
 package git
 
 import "fmt"
+import "bytes"
 
-// Status current status of a Worktree
+// Status represents the current status of a Worktree.
+// The key of the map is the path of the file.
 type Status map[string]*FileStatus
 
-func (s Status) File(filename string) *FileStatus {
-	if _, ok := (s)[filename]; !ok {
-		s[filename] = &FileStatus{}
+// File returns the FileStatus for a given path, if the FileStatus doesn't
+// exists a new FileStatus is added to the map using the path as key.
+func (s Status) File(path string) *FileStatus {
+	if _, ok := (s)[path]; !ok {
+		s[path] = &FileStatus{Worktree: Unmodified, Staging: Unmodified}
 	}
 
-	return s[filename]
-
+	return s[path]
 }
 
+// IsClean returns true if all the files aren't in Unmodified status.
 func (s Status) IsClean() bool {
 	for _, status := range s {
 		if status.Worktree != Unmodified || status.Staging != Unmodified {
@@ -25,68 +29,42 @@ func (s Status) IsClean() bool {
 }
 
 func (s Status) String() string {
-	var names []string
-	for name := range s {
-		names = append(names, name)
-	}
-
-	var output string
-	for _, name := range names {
-		status := s[name]
-		if status.Staging == 0 && status.Worktree == 0 {
+	buf := bytes.NewBuffer(nil)
+	for path, status := range s {
+		if status.Staging == Unmodified && status.Worktree == Unmodified {
 			continue
 		}
 
 		if status.Staging == Renamed {
-			name = fmt.Sprintf("%s -> %s", name, status.Extra)
+			path = fmt.Sprintf("%s -> %s", path, status.Extra)
 		}
 
-		output += fmt.Sprintf("%s%s %s\n", status.Staging, status.Worktree, name)
+		fmt.Fprintf(buf, "%c%c %s\n", status.Staging, status.Worktree, path)
 	}
 
-	return output
+	return buf.String()
 }
 
-// FileStatus status of a file in the Worktree
+// FileStatus contains the status of a file in the worktree
 type FileStatus struct {
-	Staging  StatusCode
+	// Staging is the status of a file in the staging area
+	Staging StatusCode
+	// Worktree is the status of a file in the worktree
 	Worktree StatusCode
-	Extra    string
+	// Extra contains extra information, such as the previous name in a rename
+	Extra string
 }
 
 // StatusCode status code of a file in the Worktree
-type StatusCode int8
+type StatusCode byte
 
 const (
-	Unmodified StatusCode = iota
-	Untracked
-	Modified
-	Added
-	Deleted
-	Renamed
-	Copied
-	UpdatedButUnmerged
+	Unmodified         StatusCode = ' '
+	Untracked          StatusCode = '?'
+	Modified           StatusCode = 'M'
+	Added              StatusCode = 'A'
+	Deleted            StatusCode = 'D'
+	Renamed            StatusCode = 'R'
+	Copied             StatusCode = 'C'
+	UpdatedButUnmerged StatusCode = 'U'
 )
-
-func (c StatusCode) String() string {
-	switch c {
-	case Unmodified:
-		return " "
-	case Modified:
-		return "M"
-	case Added:
-		return "A"
-	case Deleted:
-		return "D"
-	case Renamed:
-		return "R"
-	case Copied:
-		return "C"
-	case UpdatedButUnmerged:
-		return "U"
-	case Untracked:
-		return "?"
-	default:
-		return "-"
-	}
-}
diff --git a/worktree.go b/worktree.go
index f9c4ba5..ec8fad2 100644
--- a/worktree.go
+++ b/worktree.go
@@ -17,9 +17,13 @@ import (
 	"gopkg.in/src-d/go-billy.v2"
 )
 
-var ErrWorktreeNotClean = errors.New("worktree is not clean")
-var ErrSubmoduleNotFound = errors.New("submodule not found")
+var (
+	ErrWorktreeNotClean  = errors.New("worktree is not clean")
+	ErrSubmoduleNotFound = errors.New("submodule not found")
+	ErrUnstaggedChanges  = errors.New("worktree contains unstagged changes")
+)
 
+// Worktree represents a git worktree.
 type Worktree struct {
 	r  *Repository
 	fs billy.Filesystem
@@ -31,25 +35,38 @@ func (w *Worktree) Checkout(opts *CheckoutOptions) error {
 		return err
 	}
 
+	if !opts.Force {
+		unstaged, err := w.cointainsUnstagedChanges()
+		if err != nil {
+			return err
+		}
+
+		if unstaged {
+			return ErrUnstaggedChanges
+		}
+	}
+
 	c, err := w.getCommitFromCheckoutOptions(opts)
 	if err != nil {
 		return err
 	}
 
-	ro := &ResetOptions{Commit: c}
+	ro := &ResetOptions{Commit: c, Mode: MergeReset}
 	if opts.Force {
 		ro.Mode = HardReset
 	}
 
-	if err := w.Reset(ro); err != nil {
-		return err
+	if !opts.Hash.IsZero() {
+		err = w.setHEADToCommit(opts.Hash)
+	} else {
+		err = w.setHEADToBranch(opts.Branch, c)
 	}
 
-	if !opts.Hash.IsZero() {
-		return w.setCommit(opts.Hash)
+	if err != nil {
+		return err
 	}
 
-	return w.setBranch(opts.Branch, c)
+	return w.Reset(ro)
 }
 
 func (w *Worktree) getCommitFromCheckoutOptions(opts *CheckoutOptions) (plumbing.Hash, error) {
@@ -85,12 +102,12 @@ func (w *Worktree) getCommitFromCheckoutOptions(opts *CheckoutOptions) (plumbing
 	return plumbing.ZeroHash, fmt.Errorf("unsupported tag target %q", o.Type())
 }
 
-func (w *Worktree) setCommit(commit plumbing.Hash) error {
+func (w *Worktree) setHEADToCommit(commit plumbing.Hash) error {
 	head := plumbing.NewHashReference(plumbing.HEAD, commit)
 	return w.r.Storer.SetReference(head)
 }
 
-func (w *Worktree) setBranch(branch plumbing.ReferenceName, commit plumbing.Hash) error {
+func (w *Worktree) setHEADToBranch(branch plumbing.ReferenceName, commit plumbing.Hash) error {
 	target, err := w.r.Storer.Reference(branch)
 	if err != nil {
 		return err
@@ -112,6 +129,17 @@ func (w *Worktree) Reset(opts *ResetOptions) error {
 		return err
 	}
 
+	if opts.Mode == MergeReset {
+		unstaged, err := w.cointainsUnstagedChanges()
+		if err != nil {
+			return err
+		}
+
+		if unstaged {
+			return ErrUnstaggedChanges
+		}
+	}
+
 	changes, err := w.diffCommitWithStaging(opts.Commit, true)
 	if err != nil {
 		return err
@@ -133,7 +161,44 @@ func (w *Worktree) Reset(opts *ResetOptions) error {
 		}
 	}
 
-	return w.r.Storer.SetIndex(idx)
+	if err := w.r.Storer.SetIndex(idx); err != nil {
+		return err
+	}
+
+	return w.setHEADCommit(opts.Commit)
+}
+
+func (w *Worktree) cointainsUnstagedChanges() (bool, error) {
+	ch, err := w.diffStagingWithWorktree()
+	if err != nil {
+		return false, err
+	}
+
+	return len(ch) != 0, nil
+}
+
+func (w *Worktree) setHEADCommit(commit plumbing.Hash) error {
+	head, err := w.r.Reference(plumbing.HEAD, false)
+	if err != nil {
+		return err
+	}
+
+	if head.Type() == plumbing.HashReference {
+		head = plumbing.NewHashReference(plumbing.HEAD, commit)
+		return w.r.Storer.SetReference(head)
+	}
+
+	branch, err := w.r.Reference(head.Target(), false)
+	if err != nil {
+		return err
+	}
+
+	if !branch.IsBranch() {
+		return fmt.Errorf("invalid HEAD target should be a branch, found %s", branch.Type())
+	}
+
+	branch = plumbing.NewHashReference(branch.Name(), commit)
+	return w.r.Storer.SetReference(branch)
 }
 
 func (w *Worktree) checkoutChange(ch merkletrie.Change, t *object.Tree, idx *index.Index) error {
diff --git a/worktree_status.go b/worktree_status.go
index d472fde..ae7518e 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -80,17 +80,9 @@ func (w *Worktree) diffStagingWithWorktree() (merkletrie.Changes, error) {
 		return nil, err
 	}
 
-	from, err := index.NewRootNode(idx)
-	if err != nil {
-		return nil, err
-	}
-
-	to, err := filesystem.NewRootNode(w.fs)
-	if err != nil {
-		return nil, err
-	}
-
-	return merkletrie.DiffTree(from, to, IsEquals)
+	from := index.NewRootNode(idx)
+	to := filesystem.NewRootNode(w.fs)
+	return merkletrie.DiffTree(from, to, diffTreeIsEquals)
 }
 
 func (w *Worktree) diffCommitWithStaging(commit plumbing.Hash, reverse bool) (merkletrie.Changes, error) {
@@ -99,11 +91,6 @@ func (w *Worktree) diffCommitWithStaging(commit plumbing.Hash, reverse bool) (me
 		return nil, err
 	}
 
-	to, err := index.NewRootNode(idx)
-	if err != nil {
-		return nil, err
-	}
-
 	c, err := w.r.CommitObject(commit)
 	if err != nil {
 		return nil, err
@@ -114,20 +101,31 @@ func (w *Worktree) diffCommitWithStaging(commit plumbing.Hash, reverse bool) (me
 		return nil, err
 	}
 
+	to := index.NewRootNode(idx)
 	from := object.NewTreeRootNode(t)
+
 	if reverse {
-		return merkletrie.DiffTree(to, from, IsEquals)
+		return merkletrie.DiffTree(to, from, diffTreeIsEquals)
 	}
 
-	return merkletrie.DiffTree(from, to, IsEquals)
+	return merkletrie.DiffTree(from, to, diffTreeIsEquals)
 }
 
-func IsEquals(a, b noder.Hasher) bool {
-	pathA := a.(noder.Path)
-	pathB := b.(noder.Path)
-	if pathA[len(pathA)-1].IsDir() || pathB[len(pathB)-1].IsDir() {
+var emptyNoderHash = make([]byte, 24)
+
+// diffTreeIsEquals is a implementation of noder.Equals, used to compare
+// noder.Noder, it compare the content and the length of the hashes.
+//
+// Since some of the noder.Noder implementations doesn't compute a hash for
+// some directories, if any of the hashes is a 24-byte slice of zero values
+// the comparison is not done and the hashes are take as different.
+func diffTreeIsEquals(a, b noder.Hasher) bool {
+	hashA := a.Hash()
+	hashB := b.Hash()
+
+	if bytes.Equal(hashA, emptyNoderHash) || bytes.Equal(hashB, emptyNoderHash) {
 		return false
 	}
 
-	return bytes.Equal(a.Hash(), b.Hash())
+	return bytes.Equal(hashA, hashB)
 }
diff --git a/worktree_test.go b/worktree_test.go
index d32e648..59197a6 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -5,8 +5,10 @@ import (
 	"os"
 	"path/filepath"
 
+	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/filemode"
 	"gopkg.in/src-d/go-git.v4/plumbing/format/index"
+	"gopkg.in/src-d/go-git.v4/plumbing/object"
 
 	"github.com/src-d/go-git-fixtures"
 	. "gopkg.in/check.v1"
@@ -193,38 +195,42 @@ func (s *WorktreeSuite) TestCheckoutTag(c *C) {
 	c.Assert(head.Name().String(), Equals, "HEAD")
 }
 
+func (s *WorktreeSuite) TestCheckoutBisect(c *C) {
+	s.testCheckoutBisect(c, "https://github.com/src-d/go-git.git")
+}
+
+func (s *WorktreeSuite) TestCheckoutBisectSubmodules(c *C) {
+	c.Skip("not-submodule-support")
+	s.testCheckoutBisect(c, "https://github.com/git-fixtures/submodule.git")
+}
+
 // TestCheckoutBisect simulates a git bisect going through the git history and
 // checking every commit over the previous commit
-func (s *WorktreeSuite) TestCheckoutBisect(c *C) {
-	f := fixtures.ByURL("https://github.com/src-d/go-git.git").One()
-	fs := memfs.New()
+func (s *WorktreeSuite) testCheckoutBisect(c *C, url string) {
+	f := fixtures.ByURL(url).One()
 
 	w := &Worktree{
 		r:  s.NewRepository(f),
-		fs: fs,
+		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)
 
-	ref, err := w.r.Head()
-	c.Assert(err, IsNil)
-
-	commit, err := w.r.CommitObject(ref.Hash())
+	iter, err := w.r.Log(&LogOptions{})
 	c.Assert(err, IsNil)
 
-	history, err := commit.History()
-	c.Assert(err, IsNil)
-
-	for i := len(history) - 1; i >= 0; i-- {
-		err := w.Checkout(&CheckoutOptions{Hash: history[i].Hash})
+	iter.ForEach(func(commit *object.Commit) error {
+		err := w.Checkout(&CheckoutOptions{Hash: commit.Hash})
 		c.Assert(err, IsNil)
 
 		status, err := w.Status()
 		c.Assert(err, IsNil)
 		c.Assert(status.IsClean(), Equals, true)
-	}
+
+		return nil
+	})
 }
 
 func (s *WorktreeSuite) TestStatus(c *C) {
@@ -240,6 +246,84 @@ func (s *WorktreeSuite) TestStatus(c *C) {
 	c.Assert(status, HasLen, 9)
 }
 
+func (s *WorktreeSuite) TestReset(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	commit := plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9")
+
+	err := w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+
+	branch, err := w.r.Reference(plumbing.Master, false)
+	c.Assert(err, IsNil)
+	c.Assert(branch.Hash(), Not(Equals), commit)
+
+	err = w.Reset(&ResetOptions{Commit: commit})
+	c.Assert(err, IsNil)
+
+	branch, err = w.r.Reference(plumbing.Master, false)
+	c.Assert(err, IsNil)
+	c.Assert(branch.Hash(), Equals, commit)
+}
+
+func (s *WorktreeSuite) TestResetMerge(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	commit := plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9")
+
+	err := w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+
+	f, err := fs.Create(".gitignore")
+	c.Assert(err, IsNil)
+	_, err = f.Write([]byte("foo"))
+	c.Assert(err, IsNil)
+	err = f.Close()
+	c.Assert(err, IsNil)
+
+	err = w.Reset(&ResetOptions{Mode: MergeReset, Commit: commit})
+	c.Assert(err, Equals, ErrUnstaggedChanges)
+
+	branch, err := w.r.Reference(plumbing.Master, false)
+	c.Assert(err, IsNil)
+	c.Assert(branch.Hash(), Not(Equals), commit)
+}
+
+func (s *WorktreeSuite) TestResetHard(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	commit := plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9")
+
+	err := w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+
+	f, err := fs.Create(".gitignore")
+	c.Assert(err, IsNil)
+	_, err = f.Write([]byte("foo"))
+	c.Assert(err, IsNil)
+	err = f.Close()
+	c.Assert(err, IsNil)
+
+	err = w.Reset(&ResetOptions{Mode: HardReset, Commit: commit})
+	c.Assert(err, IsNil)
+
+	branch, err := w.r.Reference(plumbing.Master, false)
+	c.Assert(err, IsNil)
+	c.Assert(branch.Hash(), Equals, commit)
+}
+
 func (s *WorktreeSuite) TestStatusAfterCheckout(c *C) {
 	fs := memfs.New()
 	w := &Worktree{
@@ -253,6 +337,7 @@ func (s *WorktreeSuite) TestStatusAfterCheckout(c *C) {
 	status, err := w.Status()
 	c.Assert(err, IsNil)
 	c.Assert(status.IsClean(), Equals, true)
+
 }
 
 func (s *WorktreeSuite) TestStatusModified(c *C) {