memfs: Speed up store.put and removeAll

memfs/store.go

@@ -84,8 +84,13 @@ func (s *store) get(k string) *value {
 
 func (s *store) put(k string, v *value) *value {
 	if _, ok := s.values[k]; !ok {
-		s.keys = append(s.keys, k)
-		sort.Strings(s.keys)
+		// Insert k into s.keys at the position that keeps the slice
+		// sorted. This is O(n) for the shift, vs. O(n log n) for a
+		// full sort.Strings on every put.
+		i := sort.SearchStrings(s.keys, k)
+		s.keys = append(s.keys, "")
+		copy(s.keys[i+1:], s.keys[i:])
+		s.keys[i] = k
 	}
 
 	s.values[k] = v
@@ -109,17 +114,17 @@ func (s *store) removeAll(prefix string) {
 		return
 	}
 
-	max := len(s.keys)
-	to := -1
-	for i := from; i < max; i++ {
-		key := s.keys[i]
-		if !strings.HasPrefix(key, prefix) {
-			break
-		}
+	// Find the first key after `from` that does not start with prefix.
+	// Because s.keys is sorted, the prefix-matching range is contiguous,
+	// so we can binary-search for its end instead of scanning.
+	end := sort.Search(len(s.keys)-from, func(i int) bool {
+		return !strings.HasPrefix(s.keys[from+i], prefix)
+	}) + from
+
+	for _, key := range s.keys[from:end] {
 		delete(s.values, key)
-		to = i
 	}
-	s.keys = append(s.keys[0:from], s.keys[to+1:]...)
+	s.keys = append(s.keys[:from], s.keys[end:]...)
 }
 
 func (s *store) keyIndex(key string) int {

memfs/store_test.go

@@ -1,7 +1,9 @@
 package memfs
 
 import (
+	"fmt"
 	"io/fs"
+	"math/rand"
 	"reflect"
 	"sort"
 	"strings"
@@ -118,6 +120,26 @@ func TestStore_remove(t *testing.T) {
 	}
 }
 
+func BenchmarkStore_put(b *testing.B) {
+	// Pre-generate keys in a randomized order so that put hits the
+	// realistic insertion-into-sorted-slice path rather than always
+	// appending at the end.
+	const n = 5000
+	rng := rand.New(rand.NewSource(1))
+	keys := make([]string, n)
+	for i := range keys {
+		keys[i] = fmt.Sprintf("/dir%05d/file%05d.txt", rng.Intn(n), i)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s := newStore()
+		for _, k := range keys {
+			s.put(k, &value{name: k, mode: fs.ModePerm})
+		}
+	}
+}
+
 func TestStore_removeAll(t *testing.T) {
 	s := newStoreTest()