From ecbd067f5b1f231dd59644ca5a52188bac752b1b Mon Sep 17 00:00:00 2001
From: Tigran <hanagantig@gmail.com>
Date: Sat, 23 Sep 2023 23:02:29 +0400
Subject: [PATCH] remove mem fs with descendants

---
 memmap.go      | 20 +++++++++++++--
 memmap_test.go | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/memmap.go b/memmap.go
index d6c744e..cc539c0 100644
--- a/memmap.go
+++ b/memmap.go
@@ -285,10 +285,22 @@ func (m *MemMapFs) Remove(name string) error {
 	defer m.mu.Unlock()
 
 	if _, ok := m.getData()[name]; ok {
+		descendants := m.findDescendants(name)
+		for i := 1; i <= len(descendants); i++ {
+			descendant := descendants[len(descendants)-i]
+			descName := descendant.Name()
+			err := m.unRegisterWithParent(descName)
+			if err != nil {
+				return &os.PathError{Op: "descendant remove", Path: name, Err: err}
+			}
+			delete(m.getData(), descName)
+		}
+
 		err := m.unRegisterWithParent(name)
 		if err != nil {
 			return &os.PathError{Op: "remove", Path: name, Err: err}
 		}
+
 		delete(m.getData(), name)
 	} else {
 		return &os.PathError{Op: "remove", Path: name, Err: os.ErrNotExist}
@@ -299,14 +311,18 @@ func (m *MemMapFs) Remove(name string) error {
 func (m *MemMapFs) RemoveAll(path string) error {
 	path = normalizePath(path)
 	m.mu.Lock()
-	m.unRegisterWithParent(path)
+	_ = m.unRegisterWithParent(path)
 	m.mu.Unlock()
 
 	m.mu.RLock()
 	defer m.mu.RUnlock()
 
 	for p := range m.getData() {
-		if p == path || strings.HasPrefix(p, path+FilePathSeparator) {
+		separator := FilePathSeparator
+		if path == FilePathSeparator {
+			separator = ""
+		}
+		if p == path || strings.HasPrefix(p, path+separator) {
 			m.mu.RUnlock()
 			m.mu.Lock()
 			delete(m.getData(), p)
diff --git a/memmap_test.go b/memmap_test.go
index c47fadc..dbf89b4 100644
--- a/memmap_test.go
+++ b/memmap_test.go
@@ -918,3 +918,70 @@ func TestMemMapFsRename(t *testing.T) {
 		}
 	}
 }
+
+func TestMemMapFsRemove(t *testing.T) {
+	t.Parallel()
+
+	testData := map[string]struct {
+		dirsToCreate   []string
+		dirsToRemove   []string
+		expectedErrMsg string
+	}{
+		"Remove child before - success": {
+			dirsToCreate: []string{"/parent1/parent2/fileForDelete1.txt"},
+			dirsToRemove: []string{
+				"/parent1/parent2/fileForDelete1.txt",
+				"/parent1/parent2",
+			},
+		},
+		"Remove parent before - should return error": {
+			dirsToCreate: []string{"/parent1/parent2/fileForDelete1.txt"},
+			dirsToRemove: []string{
+				"/parent1/parent2",
+				"/parent1/parent2/fileForDelete1.txt",
+			},
+			expectedErrMsg: "remove /parent1/parent2/fileForDelete1.txt: file does not exist",
+		},
+		"Remove root and then parent1 - should return error": {
+			dirsToCreate: []string{"/root/parent1/parent2/fileForDelete1.txt"},
+			dirsToRemove: []string{
+				"/root",
+				"/root/parent1",
+			},
+			expectedErrMsg: "remove /root/parent1: file does not exist",
+		},
+		"Remove parent2 and then parent 1 - success": {
+			dirsToCreate: []string{"/parent1/parent2/fileForDelete1.txt"},
+			dirsToRemove: []string{
+				"/parent1/parent2",
+				"/parent1",
+			},
+		},
+	}
+
+	fs := &MemMapFs{}
+
+	for caseName, td := range testData {
+		_, err := fs.Stat("/")
+		if err == nil {
+			err = fs.RemoveAll("/")
+			if err != nil {
+				t.Fatalf("%s: RemoveAll %q failed: %v", fs.Name(), "/", err)
+			}
+		}
+
+		for _, toCreate := range td.dirsToCreate {
+			err = fs.MkdirAll(toCreate, os.FileMode(0775))
+			if err != nil && err.Error() != td.expectedErrMsg {
+				t.Fatalf("#CASE %v %s: Mkdir %q failed: %v", caseName, fs.Name(), toCreate, err)
+			}
+		}
+
+		for _, toRemove := range td.dirsToRemove {
+			err = fs.Remove(toRemove)
+			if err != nil && err.Error() != td.expectedErrMsg {
+				t.Fatalf("#CASE %v %s: Remove %q failed: %v", caseName, fs.Name(), toRemove, err)
+			}
+		}
+	}
+}