Co-allocate persistent vector replacements

Author: jfhamlin

internal/persistent/vector/vector.go

@@ -88,6 +88,13 @@ type TailStorage struct {
 	entry tailEntry
 }
 
+// ReplaceTailStorage reserves storage for replacing the last tail entry
+// without copying either a base tail or a delta chain.
+type ReplaceTailStorage struct {
+	entry tailEntry
+	base  tailBase
+}
+
 // Empty is an empty Vector.
 var Empty Vector = &Persistent{}
 
@@ -243,7 +250,7 @@ func (v Persistent) AssocValue(i int, val interface{}) (Persistent, bool) {
 // operation does not have to materialize and copy the tail.
 func (v Persistent) ReplaceLastValueInto(
 	val interface{},
-	storage *TailStorage,
+	storage *ReplaceTailStorage,
 ) (Persistent, bool) {
 	if v.count == 0 {
 		return Persistent{}, false
@@ -256,9 +263,11 @@ func (v Persistent) ReplaceLastValueInto(
 		v.tailDelta = &storage.entry
 		return v, true
 	}
-	tail := append([]interface{}(nil), v.baseTail()...)
-	tail[len(tail)-1] = val
-	v.tail = newTailBase(tail)
+	base := v.baseTail()
+	storage.base = tailBase(base[:len(base)-1])
+	storage.entry = tailEntry{value: val}
+	v.tail = &storage.base
+	v.tailDelta = &storage.entry
 	return v, true
 }
 

pkg/lang/vector.go

@@ -28,6 +28,11 @@ type (
 		tail vector.TailStorage
 	}
 
+	vectorReplaceStorage struct {
+		Vector
+		tail vector.ReplaceTailStorage
+	}
+
 	PersistentVector = Vector
 
 	TransientVector struct {
@@ -100,7 +105,7 @@ func (v *Vector) AssocN(i int, val any) IPersistentVector {
 
 // ReplaceLast returns a persistent vector whose final value is val.
 func (v *Vector) ReplaceLast(val any) *Vector {
-	storage := &vectorUpdateStorage{}
+	storage := &vectorReplaceStorage{}
 	result, ok := v.vec.ReplaceLastValueInto(val, &storage.tail)
 	if !ok {
 		panic("can't pop an empty vector")

pkg/lang/vector_test.go

@@ -36,23 +36,36 @@ func TestLargeVectorConsCoallocatesResultAndTail(t *testing.T) {
 }
 
 func TestVectorReplaceLastIsPersistentAndCoallocated(t *testing.T) {
-	original := NewVector()
-	for _, value := range []int{1, 2, 3, 4, 5} {
-		original = original.Cons(value).(*Vector)
-	}
-	var result *Vector
-	if got := testing.AllocsPerRun(1_000, func() {
-		result = original.ReplaceLast(9)
-	}); got != 1 {
-		t.Fatalf("replace-last allocated %v objects per call, want 1", got)
-	}
-	if got := original.Nth(4); got != 5 {
-		t.Fatalf("original final value = %v, want 5", got)
-	}
-	if got := result.Nth(4); got != 9 {
-		t.Fatalf("replacement final value = %v, want 9", got)
+	tests := []struct {
+		name     string
+		original *Vector
+	}{
+		{"base tail", NewVector(1, 2, 3, 4, 5)},
+		{"delta tail", func() *Vector {
+			vector := NewVector()
+			for _, value := range []int{1, 2, 3, 4, 5} {
+				vector = vector.Cons(value).(*Vector)
+			}
+			return vector
+		}()},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var result *Vector
+			if got := testing.AllocsPerRun(1_000, func() {
+				result = test.original.ReplaceLast(9)
+			}); got != 1 {
+				t.Fatalf("replace-last allocated %v objects per call, want 1", got)
+			}
+			if got := test.original.Nth(4); got != 5 {
+				t.Fatalf("original final value = %v, want 5", got)
+			}
+			if got := result.Nth(4); got != 9 {
+				t.Fatalf("replacement final value = %v, want 9", got)
+			}
+			runtime.KeepAlive(result)
+		})
 	}
-	runtime.KeepAlive(result)
 }
 
 func TestSmallVectorWithMetaKeepsInlineStorageAlive(t *testing.T) {