worktree: minor speedup for doAddDirectory

2 files changed, 94 insertions, 18 deletions

diff --git a/worktree_status.go b/worktree_status.go
index a26c9e5..389482d 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -281,8 +281,10 @@ func (w *Worktree) doAddDirectory(idx *index.Index, s Status, directory string,
 		}
 	}
 
+	directory = filepath.ToSlash(filepath.Clean(directory))
+
 	for name := range s {
-		if !isPathInDirectory(name, filepath.ToSlash(filepath.Clean(directory))) {
+		if !isPathInDirectory(name, directory) {
 			continue
 		}
 
@@ -301,23 +303,7 @@ func (w *Worktree) doAddDirectory(idx *index.Index, s Status, directory string,
 }
 
 func isPathInDirectory(path, directory string) bool {
-	ps := strings.Split(path, "/")
-	ds := strings.Split(directory, "/")
-
-	if len(ds) == 1 && ds[0] == "." {
-		return true
-	}
-
-	if len(ps) < len(ds) {
-		return false
-	}
-
-	for i := 0; i < len(ds); i++ {
-		if ps[i] != ds[i] {
-			return false
-		}
-	}
-	return true
+	return directory == "." || strings.HasPrefix(path, directory+"/")
 }
 
 // AddWithOptions file contents to the index,  updates the index using the
diff --git a/worktree_test.go b/worktree_test.go
index ea55ac6..24d5bd5 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -1488,6 +1488,96 @@ func (s *WorktreeSuite) TestAddRemovedInDirectory(c *C) {
 	c.Assert(file.Staging, Equals, Unmodified)
 }
 
+func (s *WorktreeSuite) TestAddRemovedInDirectoryWithTrailingSlash(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	idx, err := w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 9)
+
+	err = w.Filesystem.Remove("go/example.go")
+	c.Assert(err, IsNil)
+
+	err = w.Filesystem.Remove("json/short.json")
+	c.Assert(err, IsNil)
+
+	hash, err := w.Add("go/")
+	c.Assert(err, IsNil)
+	c.Assert(hash.IsZero(), Equals, true)
+
+	e, err := idx.Entry("go/example.go")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, plumbing.NewHash("880cd14280f4b9b6ed3986d6671f907d7cc2a198"))
+	c.Assert(e.Mode, Equals, filemode.Regular)
+
+	e, err = idx.Entry("json/short.json")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, plumbing.NewHash("c8f1d8c61f9da76f4cb49fd86322b6e685dba956"))
+	c.Assert(e.Mode, Equals, filemode.Regular)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 2)
+
+	file := status.File("go/example.go")
+	c.Assert(file.Staging, Equals, Deleted)
+
+	file = status.File("json/short.json")
+	c.Assert(file.Staging, Equals, Unmodified)
+}
+
+func (s *WorktreeSuite) TestAddRemovedInDirectoryDot(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	idx, err := w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 9)
+
+	err = w.Filesystem.Remove("go/example.go")
+	c.Assert(err, IsNil)
+
+	err = w.Filesystem.Remove("json/short.json")
+	c.Assert(err, IsNil)
+
+	hash, err := w.Add(".")
+	c.Assert(err, IsNil)
+	c.Assert(hash.IsZero(), Equals, true)
+
+	e, err := idx.Entry("go/example.go")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, plumbing.NewHash("880cd14280f4b9b6ed3986d6671f907d7cc2a198"))
+	c.Assert(e.Mode, Equals, filemode.Regular)
+
+	e, err = idx.Entry("json/short.json")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, plumbing.NewHash("c8f1d8c61f9da76f4cb49fd86322b6e685dba956"))
+	c.Assert(e.Mode, Equals, filemode.Regular)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 2)
+
+	file := status.File("go/example.go")
+	c.Assert(file.Staging, Equals, Deleted)
+
+	file = status.File("json/short.json")
+	c.Assert(file.Staging, Equals, Deleted)
+}
+
 func (s *WorktreeSuite) TestAddSymlink(c *C) {
 	dir, clean := s.TemporalDir()
 	defer clean()