From 34fac9d5d0d7ce09649d6094074273a9cbfe16f2 Mon Sep 17 00:00:00 2001
From: Tigran <hanagantig@gmail.com>
Date: Fri, 8 Sep 2023 01:40:48 +0400
Subject: [PATCH] remove mem fs with descendants

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

diff --git a/memmap.go b/memmap.go
index e6b7d70..574d971 100644
--- a/memmap.go
+++ b/memmap.go
@@ -16,9 +16,12 @@ package afero
 import (
 	"fmt"
 	"io"
+
 	"log"
 	"os"
 	"path/filepath"
+
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -88,6 +91,24 @@ func (m *MemMapFs) findParent(f *mem.FileData) *mem.FileData {
 	return pfile
 }
 
+func (m *MemMapFs) findDescendants(name string) []*mem.FileData {
+	fData := m.getData()
+	descendants := make([]*mem.FileData, 0, len(fData))
+	for p, dFile := range fData {
+		if strings.HasPrefix(p, name+FilePathSeparator) {
+			descendants = append(descendants, dFile)
+		}
+	}
+
+	sort.Slice(descendants, func(i, j int) bool {
+		cur := len(strings.Split(descendants[i].Name(), FilePathSeparator))
+		next := len(strings.Split(descendants[j].Name(), FilePathSeparator))
+		return cur < next
+	})
+
+	return descendants
+}
+
 func (m *MemMapFs) registerWithParent(f *mem.FileData, perm os.FileMode) {
 	if f == nil {
 		return
@@ -264,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}
@@ -278,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 52a492e..9d68b9a 100644
--- a/memmap_test.go
+++ b/memmap_test.go
@@ -833,3 +833,70 @@ func TestMemFsRenameDir(t *testing.T) {
 		t.Errorf("Cannot recreate the subdir in the source dir: %s", err)
 	}
 }
+
+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)
+			}
+		}
+	}
+}