ollama/x/mlxrunner/prefix_cache_test.go

package mlxrunner

import (
	"slices"
	"testing"
	"time"

	"github.com/ollama/ollama/x/mlxrunner/cache"
	"github.com/ollama/ollama/x/mlxrunner/mlx"
	"github.com/ollama/ollama/x/mlxrunner/model/base"
)

// snapshotTracker records every fakeSnapshot created and every Close() call
// so tests can detect leaked (created but never closed) or double-closed snapshots.
type snapshotTracker struct {
	all []*fakeSnapshot
}

func (tr *snapshotTracker) track(s *fakeSnapshot) {
	if s == nil {
		return
	}
	s.tracker = tr
	tr.all = append(tr.all, s)
}

// Fake caches that store actual token sequences so tests can verify the right
// data was restored, not just the right offset.

// fakePending mirrors the production pendingSnapshots capture machinery for the
// fakes: it schedules offsets, captures as feed crosses each one (edge-local,
// via a running base cursor), and returns the captures in scheduled order.
// capture is the owning fake's snapshot-at-an-offset function.
type fakePending struct {
	offsets  []int
	captured []cache.Snapshot
	base     int
}

func (p *fakePending) prepare(currentOffset int, offsets []int) {
	p.offsets = slices.Clone(offsets)
	p.captured = make([]cache.Snapshot, len(offsets))
	p.base = currentOffset
}

func (p *fakePending) take() []cache.Snapshot {
	out := p.captured
	p.offsets, p.captured = nil, nil
	return out
}

// feedCapturing advances the cache from start over the fed tokens, capturing at
// every scheduled offset crossed (including start and end). capture(from,
// reached) produces the snapshot for the edge ending at reached;
// advance(tokens) appends a token segment to the live state. Segment boundaries
// fall at scheduled offsets.
func (p *fakePending) feedCapturing(start int, tokens []int32, capture func(from, reached int) cache.Snapshot, advance func([]int32)) {
	end := start + len(tokens)
	captureAt := func(reached int) {
		fired := false
		for i, o := range p.offsets {
			if p.captured[i] == nil && o == reached {
				p.captured[i] = capture(p.base, reached)
				fired = true
			}
		}
		// The base advances only when a capture fires, mirroring
		// pendingSnapshots.captureReached.
		if fired {
			p.base = reached
		}
	}

	if len(p.offsets) == 0 {
		advance(tokens)
		return
	}

	captureAt(start)
	prev := 0
	for cut := start + 1; cut < end; cut++ {
		if !slices.Contains(p.offsets, cut) {
			continue
		}
		advance(tokens[prev : cut-start])
		captureAt(cut)
		prev = cut - start
	}
	advance(tokens[prev:])
	captureAt(end)
}

// fakeSnapshot stores a copy of the token sub-sequence it covers.
type fakeSnapshot struct {
	tokens   []int32
	from, to int
	byteSize int // configurable for eviction tests

	tracker    *snapshotTracker
	closeCount int

	onMaterialize func(delta int)
}

func (s *fakeSnapshot) Size() int                             { return s.byteSize }
func (s *fakeSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
func (s *fakeSnapshot) Close()                                { s.closeCount++ }

// materialize simulates a lazy snapshot copying out: grow byteSize by delta
// and fire the trie's hook if one is attached. Tests use this to verify the
// trie's counter responds to materialization events.
func (s *fakeSnapshot) materialize(delta int) {
	s.byteSize += delta
	if s.onMaterialize != nil {
		s.onMaterialize(delta)
		s.onMaterialize = nil
	}
}

// fakeRewindableCache tracks the full token sequence and supports
// arbitrary rewind via Restore(nil, target).
type fakeRewindableCache struct {
	tokens  []int32
	tracker *snapshotTracker
	pending fakePending
}

func (c *fakeRewindableCache) feed(tokens []int32) {
	c.pending.feedCapturing(len(c.tokens), tokens,
		func(from, reached int) cache.Snapshot { return c.Snapshot(from) },
		func(seg []int32) { c.tokens = append(c.tokens, seg...) })
}

func (c *fakeRewindableCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
	return nil, nil
}
func (c *fakeRewindableCache) State() []*mlx.Array { return nil }
func (c *fakeRewindableCache) Offset() int         { return len(c.tokens) }

func (c *fakeRewindableCache) Free() {
	c.tokens = nil
}

func (c *fakeRewindableCache) Snapshot(fromOffset int) cache.Snapshot {
	if fromOffset >= len(c.tokens) {
		return nil
	}
	from := fromOffset
	if from < 0 {
		from = 0
	}
	s := &fakeSnapshot{
		tokens: slices.Clone(c.tokens[from:]),
		from:   from,
		to:     len(c.tokens),
	}
	c.tracker.track(s)
	return s
}

func (c *fakeRewindableCache) Restore(snapshot cache.Snapshot, target int) bool {
	if target < 0 {
		return false
	}

	if snapshot == nil {
		if target > len(c.tokens) {
			return false
		}
		c.tokens = c.tokens[:target]
		return true
	}
	s := snapshot.(*fakeSnapshot)
	if target > s.to || len(c.tokens) < s.from {
		return false
	}
	c.tokens = append(c.tokens[:s.from], s.tokens...)
	if target < len(c.tokens) {
		c.tokens = c.tokens[:target]
	}
	return true
}

func (c *fakeRewindableCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
	if parent == nil || child == nil {
		if parent != nil {
			parent.Close()
		}
		if child != nil {
			child.Close()
		}
		return nil
	}
	p := parent.(*fakeSnapshot)
	ch := child.(*fakeSnapshot)
	merged := make([]int32, len(p.tokens)+len(ch.tokens))
	copy(merged, p.tokens)
	copy(merged[len(p.tokens):], ch.tokens)
	s := &fakeSnapshot{
		tokens:   merged,
		from:     p.from,
		to:       ch.to,
		byteSize: p.byteSize + ch.byteSize,
	}
	c.tracker.track(s)
	p.Close()
	ch.Close()
	return s
}

func (c *fakeRewindableCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
	if snapshot == nil {
		return nil, nil
	}
	s := snapshot.(*fakeSnapshot)
	relAt := at - s.from
	if relAt <= 0 {
		return nil, snapshot
	}
	if relAt >= len(s.tokens) {
		return snapshot, nil
	}
	p := &fakeSnapshot{
		tokens:   slices.Clone(s.tokens[:relAt]),
		from:     s.from,
		to:       at,
		byteSize: s.byteSize,
	}
	ch := &fakeSnapshot{
		tokens:   slices.Clone(s.tokens[relAt:]),
		from:     at,
		to:       s.to,
		byteSize: s.byteSize,
	}
	c.tracker.track(p)
	c.tracker.track(ch)
	s.Close()
	return p, ch
}

func (c *fakeRewindableCache) PrepareSnapshots(offsets []int) {
	c.pending.prepare(len(c.tokens), offsets)
}
func (c *fakeRewindableCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }

// fakeSlidingWindowCache models RotatingKVCache semantics: stores the full
// token sequence but only the trailing maxSize tokens are "live" in the window.
// Once the window fills, live rewind is impossible without a snapshot.
type fakeSlidingWindowCache struct {
	tokens  []int32
	maxSize int
	tracker *snapshotTracker
	pending fakePending
}

func (c *fakeSlidingWindowCache) feed(tokens []int32) {
	c.pending.feedCapturing(len(c.tokens), tokens,
		func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
		func(seg []int32) { c.tokens = append(c.tokens, seg...) })
}

func (c *fakeSlidingWindowCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
	return nil, nil
}
func (c *fakeSlidingWindowCache) State() []*mlx.Array { return nil }
func (c *fakeSlidingWindowCache) Offset() int         { return len(c.tokens) }

func (c *fakeSlidingWindowCache) Free() {
	c.tokens = nil
}

func (c *fakeSlidingWindowCache) Snapshot(fromOffset int) cache.Snapshot {
	if len(c.tokens) == 0 || len(c.tokens) <= fromOffset {
		return nil
	}
	// Snapshot captures the full window state (like RotatingKVCache.Snapshot).
	s := &fakeSnapshot{
		tokens: slices.Clone(c.tokens),
		from:   0,
		to:     len(c.tokens),
	}
	c.tracker.track(s)
	return s
}

func (c *fakeSlidingWindowCache) Restore(snapshot cache.Snapshot, target int) bool {
	if target < 0 {
		return false
	}

	if snapshot == nil {
		if target >= len(c.tokens) {
			return target == len(c.tokens)
		}
		// Live rewind only works when buffer hasn't filled (offset <= maxSize).
		if len(c.tokens) > c.maxSize {
			return false
		}
		c.tokens = c.tokens[:target]
		return true
	}
	s := snapshot.(*fakeSnapshot)
	if target > s.to {
		return false
	}
	// Reject if clamping would leave an incomplete window
	// (matches RotatingKVCache behavior).
	if target < s.to && s.to > c.maxSize {
		return false
	}
	c.tokens = slices.Clone(s.tokens)
	if target < len(c.tokens) {
		c.tokens = c.tokens[:target]
	}
	return true
}

func (c *fakeSlidingWindowCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
	// Child supersedes parent for sliding window (full window state).
	if parent != nil {
		parent.Close()
	}
	return child
}

func (c *fakeSlidingWindowCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
	// Can't split a ring buffer at an arbitrary point.
	return nil, snapshot
}

func (c *fakeSlidingWindowCache) PrepareSnapshots(offsets []int) {
	c.pending.prepare(len(c.tokens), offsets)
}
func (c *fakeSlidingWindowCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }

// fakeRecurrentCache models RecurrentCache semantics: stores tokens
// but cannot rewind without a snapshot.
type fakeRecurrentCache struct {
	tokens  []int32
	tracker *snapshotTracker
	pending fakePending
}

func (c *fakeRecurrentCache) feed(tokens []int32) {
	c.pending.feedCapturing(len(c.tokens), tokens,
		func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
		func(seg []int32) { c.tokens = append(c.tokens, seg...) })
}

func (c *fakeRecurrentCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
	return nil, nil
}
func (c *fakeRecurrentCache) State() []*mlx.Array { return nil }
func (c *fakeRecurrentCache) Offset() int         { return len(c.tokens) }

func (c *fakeRecurrentCache) Free() {
	c.tokens = nil
}

func (c *fakeRecurrentCache) Snapshot(fromOffset int) cache.Snapshot {
	// Recurrent state is cumulative; snapshot captures the full state.
	if len(c.tokens) == 0 {
		return nil
	}
	s := &fakeSnapshot{
		tokens: slices.Clone(c.tokens),
		from:   0,
		to:     len(c.tokens),
	}
	c.tracker.track(s)
	return s
}

func (c *fakeRecurrentCache) Restore(snapshot cache.Snapshot, target int) bool {
	if snapshot == nil {
		return target == len(c.tokens) // can only no-op
	}
	s := snapshot.(*fakeSnapshot)
	if target != s.to {
		return false // cumulative state requires exact match
	}
	c.tokens = slices.Clone(s.tokens)
	return true
}

func (c *fakeRecurrentCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
	// Child supersedes parent for cumulative state.
	if parent != nil {
		parent.Close()
	}
	return child
}

func (c *fakeRecurrentCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
	return nil, snapshot // can't split cumulative state
}

func (c *fakeRecurrentCache) PrepareSnapshots(offsets []int) {
	c.pending.prepare(len(c.tokens), offsets)
}
func (c *fakeRecurrentCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }

type feedableCache interface {
	cache.Cache
	feed(tokens []int32)
}

// testEnv encapsulates a prefixCache and its fake caches for a test scenario.
type testEnv struct {
	pc         *prefixCache
	caches     []cache.Cache // typed references for assertions
	tracker    *snapshotTracker
	rewindable bool // true when all caches support arbitrary Restore(nil, target)
}

// newTransformerEnv creates a test environment with a single rewindable cache
// (pure transformer model).
func newTransformerEnv() *testEnv {
	tracker := &snapshotTracker{}
	caches := []cache.Cache{&fakeRewindableCache{tracker: tracker}}
	return &testEnv{
		pc:         &prefixCache{caches: caches},
		caches:     caches,
		tracker:    tracker,
		rewindable: true,
	}
}

// newSlidingWindowEnv creates a test environment with one rewindable cache and
// one sliding window cache (Mistral-style architecture). The sliding window
// maxSize is set small enough that test sequences fill it, making
// Restore(nil, target) fail — the same behavior as production models where
// the window fills after a few turns.
func newSlidingWindowEnv() *testEnv {
	tr := &snapshotTracker{}
	rc := &fakeRewindableCache{tracker: tr}
	sw := &fakeSlidingWindowCache{maxSize: 4, tracker: tr}
	caches := []cache.Cache{rc, sw}
	return &testEnv{
		pc:         &prefixCache{caches: caches},
		caches:     caches,
		tracker:    tr,
		rewindable: false,
	}
}

// newStatelessLayerEnv gives the cache slice nil holes, the layout a hybrid
// model produces when some layers own no state.
func newStatelessLayerEnv() *testEnv {
	tr := &snapshotTracker{}
	rc := &fakeRewindableCache{tracker: tr}
	nrc := &fakeRecurrentCache{tracker: tr}
	caches := []cache.Cache{nrc, nil, rc, nil}
	return &testEnv{
		pc:         &prefixCache{caches: caches},
		caches:     caches,
		tracker:    tr,
		rewindable: false,
	}
}

// newRecurrentEnv creates a test environment with one rewindable cache and one
// non-rewindable cache (Jamba-style architecture).
func newRecurrentEnv() *testEnv {
	tr := &snapshotTracker{}
	rc := &fakeRewindableCache{tracker: tr}
	nrc := &fakeRecurrentCache{tracker: tr}
	caches := []cache.Cache{rc, nrc}
	return &testEnv{
		pc:         &prefixCache{caches: caches},
		caches:     caches,
		tracker:    tr,
		rewindable: false,
	}
}

// assertAllTokens checks that every cache in the environment contains exactly
// the expected token sequence.
func (e *testEnv) assertAllTokens(t *testing.T, label string, expected []int32) {
	t.Helper()
	var first cache.Cache
	for i, c := range e.caches {
		if c == nil {
			continue
		}
		assertTokens(t, label, c, expected)
		// Verify all caches report the same offset.
		if first == nil {
			first = c
			continue
		}
		if c.Offset() != first.Offset() {
			t.Errorf("%s: cache %d offset=%d != first cache offset=%d",
				label, i, c.Offset(), first.Offset())
		}
	}
}

// simulateRequest mirrors the production pipeline lifecycle:
//   begin -> schedule snapshots -> prefill in one pass -> attach snapshots -> generate -> close

type requestResult struct {
	remaining        []int32
	pendingSnapshots int
}

// simulateRequest runs a request through the harness. If userSnapshotAt > 0,
// a user snapshot is requested at that offset during prefill.
func simulateRequest(t *testing.T, pc *prefixCache, inputs, generated []int32, userSnapshotAt ...int) requestResult {
	t.Helper()

	session := pc.begin(inputs, nil)
	var snapshotOffsets []int
	for _, at := range userSnapshotAt {
		if at > 0 {
			snapshotOffsets = append(snapshotOffsets, at)
		}
	}

	result := requestResult{
		remaining:        slices.Clone(session.remaining),
		pendingSnapshots: len(session.pendingSnapshots),
	}

	assertCacheOffsetAlignment(t, pc, "after begin")

	baseOffset := pc.minCacheOffset()
	seed := len(inputs) - 1

	// Prefill: schedule the pending snapshots, feed the prompt up to the
	// seed token in one pass (the caches self-segment at the scheduled
	// offsets), then attach the captures to the trie.
	session.schedulePrefillSnapshots(snapshotOffsets)
	if baseOffset < seed {
		feedAll(pc.caches, inputs[baseOffset:seed])
	}
	session.attachPrefillSnapshots()

	assertCacheOffsetAlignment(t, pc, "after prefill")

	// Decode: feed the seed and all but the last generated token.
	if len(generated) > 0 {
		session.outputs = generated
		feedAll(pc.caches, inputs[seed:])
		feedAll(pc.caches, generated[:len(generated)-1])
	}

	assertCacheOffsetAlignment(t, pc, "before close")
	session.close()
	return result
}

func feedAll(caches []cache.Cache, tokens []int32) {
	for _, c := range caches {
		if fc, ok := c.(feedableCache); ok {
			fc.feed(tokens)
		}
	}
}

// assertCacheOffsetAlignment verifies all caches report the same offset.
func assertCacheOffsetAlignment(t *testing.T, pc *prefixCache, label string) {
	t.Helper()
	expected := -1
	for i, c := range pc.caches {
		if c == nil {
			continue
		}
		if expected < 0 {
			expected = c.Offset()
			continue
		}
		if got := c.Offset(); got != expected {
			t.Errorf("%s: cache %d offset=%d != first cache offset=%d", label, i, got, expected)
		}
	}
}

// assertTokens checks that a feedable cache contains the expected token sequence.
// For sliding window caches, only the trailing maxSize tokens are checked.
func assertTokens(t *testing.T, label string, c cache.Cache, expected []int32) {
	t.Helper()
	switch fc := c.(type) {
	case *fakeRewindableCache:
		if !slices.Equal(fc.tokens, expected) {
			t.Errorf("%s: rewindable tokens = %v, want %v", label, fc.tokens, expected)
		}
	case *fakeSlidingWindowCache:
		// Sliding window stores full history but only trailing maxSize are live.
		// Verify the full token sequence matches (the window semantics are
		// enforced by Snapshot/Restore, not by the token log).
		if !slices.Equal(fc.tokens, expected) {
			t.Errorf("%s: sliding window tokens = %v, want %v", label, fc.tokens, expected)
		}
	case *fakeRecurrentCache:
		if !slices.Equal(fc.tokens, expected) {
			t.Errorf("%s: non-rewindable tokens = %v, want %v", label, fc.tokens, expected)
		}
	default:
		t.Fatalf("%s: unknown cache type %T", label, c)
	}
}

// checkTrieInvariants walks the trie and checks structural invariants.
func checkTrieInvariants(t *testing.T, root *trieNode) {
	t.Helper()
	walkNodes(root, func(n *trieNode) bool {
		if n.parent != nil {
			if n.startOffset() != n.parent.endOffset {
				t.Errorf("node [%d,%d): startOffset %d != parent endOffset %d",
					n.startOffset(), n.endOffset, n.startOffset(), n.parent.endOffset)
			}
		}
		if len(n.tokens) != n.endOffset-n.startOffset() {
			t.Errorf("node [%d,%d): token count %d != offset span %d",
				n.startOffset(), n.endOffset, len(n.tokens), n.endOffset-n.startOffset())
		}
		for _, c := range n.children {
			if c.parent != n {
				t.Errorf("child [%d,%d) parent mismatch", c.startOffset(), c.endOffset)
			}
		}
		// No two siblings should start with the same token.
		seen := make(map[trieKey]bool)
		for _, c := range n.children {
			if len(c.tokens) > 0 {
				first := c.tokens[0]
				if seen[first] {
					t.Errorf("node [%d,%d): duplicate sibling first token %d",
						n.startOffset(), n.endOffset, first)
				}
				seen[first] = true
			}
		}
		return true
	})
}

// checkSnapshotLeaks verifies that every tracked snapshot is either still live
// in the trie (closeCount == 0) or has been closed exactly once. It reports
// leaked snapshots (not in trie, never closed) and double-closes.
func checkSnapshotLeaks(t *testing.T, tracker *snapshotTracker, root *trieNode) {
	t.Helper()
	if tracker == nil {
		return
	}

	// Collect all live snapshots still referenced by trie nodes.
	live := make(map[*fakeSnapshot]bool)
	walkNodes(root, func(n *trieNode) bool {
		for _, s := range n.snapshots {
			if s != nil {
				if fs, ok := s.(*fakeSnapshot); ok {
					live[fs] = true
				}
			}
		}
		return true
	})

	for i, s := range tracker.all {
		if live[s] {
			if s.closeCount != 0 {
				t.Errorf("snapshot #%d [%d,%d) is still in trie but was closed %d time(s)",
					i, s.from, s.to, s.closeCount)
			}
		} else {
			if s.closeCount == 0 {
				t.Errorf("snapshot #%d [%d,%d) leaked: created but never closed and not in trie",
					i, s.from, s.to)
			} else if s.closeCount > 1 {
				t.Errorf("snapshot #%d [%d,%d) double-closed: closed %d times",
					i, s.from, s.to, s.closeCount)
			}
		}
	}
}

// forEachEnv runs fn as subtests for three realistic model configurations —
// pure transformer, transformer + sliding window (Mistral-style), and
// transformer + recurrent (Jamba-style) — each with and without a draft
// look-ahead. Leak checking runs automatically at the end of each subtest.
func forEachEnv(t *testing.T, fn func(t *testing.T, env *testEnv)) {
	t.Helper()
	envs := []struct {
		name string
		make func() *testEnv
	}{
		{"Transformer", newTransformerEnv},
		{"SlidingWindow", newSlidingWindowEnv},
		{"Recurrent", newRecurrentEnv},
		{"StatelessLayers", newStatelessLayerEnv},
	}
	for _, e := range envs {
		for _, lookahead := range []int{0, 1} {
			name := e.name
			if lookahead > 0 {
				name += "+Lookahead"
			}
			t.Run(name, func(t *testing.T) {
				env := e.make()
				env.pc.draftLookahead = lookahead
				t.Cleanup(func() {
					checkSnapshotLeaks(t, env.tracker, env.pc.root)
				})
				fn(t, env)
			})
		}
	}
}

// TestBranchCreationAndReuse exercises the core multi-conversation lifecycle:
// two conversations share a prefix and diverge, creating a branch point.
// A third conversation extends the first. Verifies trie structure, cache
// hit lengths, and that semantic caches contain the correct token sequences.
func TestBranchCreationAndReuse(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc

		// Request A: [1,2,3,4,5,6,7,8] + generate [20,21] — full miss.
		resA := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8}, []int32{20, 21})
		if len(resA.remaining) != 8 {
			t.Fatalf("A: remaining = %d, want 8 (full miss)", len(resA.remaining))
		}
		env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 6, 7, 8, 20})

		// Verify trie was populated by close(): everything in the caches
		// is findable; the last generated token is not.
		seqA := []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21}
		_, mA := findBestMatch(pc.root, pc.key(effectiveKeyTokens(seqA, nil)))
		if want := len(seqA) - 1; mA != want {
			t.Fatalf("A findable: expected %d matched, got %d", want, mA)
		}

		// Request B: [1,2,3,4,5,10,11,12] — shares 5-token prefix with A.
		// For rewindable caches, switchToPath rewinds to the match point
		// (one below the shared tokens with a look-ahead) so only the suffix
		// needs evaluation. For non-rewindable caches (RecurrentCache), the
		// rewind fails and freeAll fires.
		resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12}, []int32{30, 31})
		if env.rewindable {
			if resB.pendingSnapshots != 0 {
				t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
			}
			if want := 3 + pc.draftLookahead; len(resB.remaining) != want {
				t.Fatalf("B: remaining = %d, want %d (rewind to match point)", len(resB.remaining), want)
			}
		} else {
			if resB.pendingSnapshots != 1 {
				t.Fatalf("B: pendingSnapshots = %d, want 1", resB.pendingSnapshots)
			}
			if len(resB.remaining) != 8 {
				t.Fatalf("B: remaining = %d, want 8 (freeAll fallback)", len(resB.remaining))
			}
		}
		env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 12, 30})

		// Both A and B should be findable in the trie.
		_, mA2 := findBestMatch(pc.root, pc.key(effectiveKeyTokens(seqA, nil)))
		if mA2 < 5 {
			t.Fatalf("A still findable: expected >= 5 matched, got %d", mA2)
		}
		_, mB := findBestMatch(pc.root, pc.key(effectiveKeyTokens([]int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31}, nil)))
		if mB < 5 {
			t.Fatalf("B findable: expected >= 5 matched, got %d", mB)
		}

		// Request C: [1,2,3,4,5,6,7,8,40,41] — extends A's prefix.
		// Should get a cache hit for the shared prefix.
		resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8, 40, 41}, nil)
		if len(resC.remaining) >= 10 {
			t.Fatalf("C: remaining = %d, want < 10 (should get cache hit)", len(resC.remaining))
		}
		env.assertAllTokens(t, "after C", []int32{1, 2, 3, 4, 5, 6, 7, 8, 40})

		checkTrieInvariants(t, pc.root)
	})
}

// TestExactMatchSeedBehavior verifies the holdback mechanism: when the exact
// same prompt is requested twice, the cache does not overclaim cached work.
// The last token must be re-evaluated to seed generation.
func TestExactMatchSeedBehavior(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc

		// Request A: first time.
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})

		// Request B: identical prompt. Holdback means matched=4, partial in
		// the 5-token edge. For rewindable caches, switchToPath rewinds to
		// offset 4, so only the held-back token needs re-evaluation. For
		// non-rewindable caches, the rewind fails and freeAll fires.
		resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{20, 21})
		if env.rewindable {
			if len(resB.remaining) != 1 {
				t.Fatalf("B: remaining = %d, want 1 (rewind to holdback point)", len(resB.remaining))
			}
			if resB.pendingSnapshots != 0 {
				t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
			}
		} else {
			if len(resB.remaining) != 5 {
				t.Fatalf("B: remaining = %d, want 5 (freeAll fallback)", len(resB.remaining))
			}
			if resB.pendingSnapshots != 1 {
				t.Fatalf("B: pendingSnapshots = %d, want 1", resB.pendingSnapshots)
			}
		}
		env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 20})

		checkTrieInvariants(t, pc.root)
	})
}

// TestConversationResumption tests the most common pattern: user sends a message,
// gets a response, then sends a follow-up. The follow-up should reuse the cached
// prefix (system prompt + first turn + assistant response).
func TestConversationResumption(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc

		// Turn 1: system prompt + user message, assistant generates response.
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11, 12})
		env.assertAllTokens(t, "turn 1", []int32{1, 2, 3, 4, 5, 10, 11})

		// Turn 2: full history + new user message; the match lands exactly
		// at the caches' offset, so even exact-offset state is reused.
		resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21}, []int32{30})
		if len(resB.remaining) > 5 {
			t.Fatalf("turn 2: remaining = %d, want <= 5 (should reuse most of history)", len(resB.remaining))
		}
		env.assertAllTokens(t, "turn 2", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21})

		// Turn 3: even longer history.
		resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41}, nil)
		if len(resC.remaining) > 5 {
			t.Fatalf("turn 3: remaining = %d, want <= 5", len(resC.remaining))
		}
		env.assertAllTokens(t, "turn 3", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40})

		checkTrieInvariants(t, pc.root)
	})
}

// TestEvictionPreservesActiveConversations creates multiple conversations sharing
// a system prompt, triggers eviction via large snapshot sizes, and verifies the
// active path and shared prefix survive while memory stays bounded.
func TestEvictionPreservesActiveConversations(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		systemPrompt := []int32{1, 2, 3, 4, 5}

		// Create 5 conversations with unique suffixes.
		for i := range 5 {
			suffix := []int32{int32(100 + i*10), int32(101 + i*10), int32(102 + i*10)}
			inputs := append(slices.Clone(systemPrompt), suffix...)
			simulateRequest(t, pc, inputs, []int32{int32(200 + i)})
		}

		// Inflate snapshot sizes to trigger eviction.
		walkNodes(pc.root, func(n *trieNode) bool {
			if !n.hasSnapshots() {
				return true
			}
			snaps := make([]cache.Snapshot, len(n.snapshots))
			for i, s := range n.snapshots {
				if s != nil {
					snaps[i] = &fakeSnapshot{byteSize: 2 * 1024 * 1024 * 1024} // 2 GiB per snapshot
				}
			}
			n.setSnapshots(snaps, &pc.pagedOutBytes)
			return true
		})

		// Run eviction.
		pc.enforceEvictionPolicy()

		// Memory should be within limits.
		if pc.pagedOutBytes > maxPagedOutBytes {
			t.Fatalf("pagedOutBytes = %d, want <= %d", pc.pagedOutBytes, maxPagedOutBytes)
		}

		// The branch point and the frontier survive.
		if len(pc.activePath) < 2 {
			t.Fatalf("activePath should have >= 2 nodes, got %d", len(pc.activePath))
		}

		// System prompt prefix should still be findable (multi-child
		// branch points are protected from eviction entirely).
		_, matched := findBestMatch(pc.root, pc.key(effectiveKeyTokens(systemPrompt, nil)))
		if want := len(pc.key(effectiveKeyTokens(systemPrompt, nil))); matched < want {
			t.Fatalf("system prompt match = %d, want %d", matched, want)
		}

		checkTrieInvariants(t, pc.root)
	})
}

// TestUserSnapshotPreservesRestorePoint verifies that user-created snapshots
// (snapshot(true)) are exact restore points that resist structural changes:
//   - A user node forces new tokens into a child instead of extending in-place
//   - A prompt diverging at the snapshot resumes there, even for caches that
//     cannot rewind
//   - The snapshot remains restorable after other branches are added
func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		inputs := []int32{1, 2, 3, 4, 5}

		// Request A: user snapshot at offset 4, then generate.
		simulateRequest(t, pc, inputs, []int32{10, 11}, 4)

		assertUserNodeExists(t, pc, "after A")

		// Request B: extends A's prefix. The user node should force tokens
		// into a child rather than extending in-place.
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21}, nil)
		env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 20})
		assertUserNodeExists(t, pc, "after B")

		// Request C: diverge at the snapshot — prefill resumes at its capture
		// point, one token lower with a look-ahead (the boundary token
		// re-evaluates to rebuild its draft pair).
		divergeC := append(slices.Clone(inputs[:4]), 30, 31)
		resC := simulateRequest(t, pc, divergeC, []int32{40})
		if want := divergeC[4-pc.draftLookahead:]; !slices.Equal(resC.remaining, want) {
			t.Fatalf("C: remaining = %v, want %v", resC.remaining, want)
		}

		// Request D: switch back to A's branch — user snapshot still restorable.
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21, 50}, nil)
		env.assertAllTokens(t, "back to A", []int32{1, 2, 3, 4, 5, 10, 11, 20, 21})

		checkTrieInvariants(t, pc.root)
	})
}

// TestUserSnapshotResistsAutoMerge verifies that when a sibling leaf is evicted,
// a user-marked parent node is not auto-merged with its remaining single child.
func TestUserSnapshotResistsAutoMerge(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		inputs := []int32{1, 2, 3, 4, 5}

		// Request A: user snapshot at offset 3, then continue to offset 5.
		simulateRequest(t, pc, inputs, []int32{10}, 3)

		// Request B: diverges at the user node, creating a second child.
		simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20})

		userNode := findUserNode(t, pc)
		if len(userNode.children) != 2 {
			t.Fatalf("user node children = %d, want 2", len(userNode.children))
		}

		// Inflate snapshot sizes so that evicting the non-active branch alone
		// brings the trie under budget, leaving the user node with one child.
		var kept, evicted int
		walkNodes(pc.root, func(n *trieNode) bool {
			for _, s := range n.snapshots {
				if s == nil {
					continue
				}
				if n.parent == userNode && !slices.Contains(pc.activePath, n) {
					evicted++
				} else {
					kept++
				}
			}
			return true
		})
		if evicted == 0 {
			t.Fatal("no snapshots on the non-active branch")
		}
		size := int(maxPagedOutBytes) / kept
		walkNodes(pc.root, func(n *trieNode) bool {
			if !n.hasSnapshots() {
				return true
			}
			snaps := make([]cache.Snapshot, len(n.snapshots))
			for i, s := range n.snapshots {
				if s != nil {
					snaps[i] = &fakeSnapshot{byteSize: size}
				}
			}
			n.setSnapshots(snaps, &pc.pagedOutBytes)
			return true
		})
		pc.enforceEvictionPolicy()

		// The user node should still exist (not auto-merged) even with one child.
		assertUserNodeExists(t, pc, "after eviction")

		checkTrieInvariants(t, pc.root)
	})
}

// TestSnapshotBeyondPrefillSkipped verifies that a snapshot scheduled at an
// offset the prefill never reaches (prefill leaves one token for decode
// seeding, so the last token is never written during prefill) is dropped rather
// than materialized as a trie node claiming tokens the cache never wrote.
func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		inputs := []int32{1, 2, 3, 4, 5}

		session := pc.begin(inputs, nil)
		// Request a snapshot at 3 and one at len(inputs); captures land at
		// the requested prefix minus the look-ahead.
		session.schedulePrefillSnapshots([]int{3, len(inputs)})
		// Prefill writes all but the final token (mirrors total-processed > 1).
		feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1])
		session.attachPrefillSnapshots()

		// The first request became a node at its capture point; nothing may
		// claim offsets the prefill never wrote.
		if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) {
			t.Errorf("no trie node at capture point %d", at)
		}
		reached := pc.minCacheOffset()
		walkNodes(pc.root, func(n *trieNode) bool {
			if n.endOffset > reached {
				t.Errorf("trie node materialized at unwritten offset %d", n.endOffset)
			}
			return true
		})

		checkTrieInvariants(t, pc.root)
	})
}

// TestAtomicMediaBoundaries verifies that a non-causal media item is never
// split by the cache: a capture scheduled inside its tokens lands at their
// end, and a prompt whose match ends inside them resumes before them.
func TestAtomicMediaBoundaries(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		inputs := []int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
		const itemPos, itemLen = 3, 6
		items := []mediaItem{{pos: itemPos, length: itemLen, fold: 1 << 31, item: &base.PreparedItem{}}}

		session := pc.begin(inputs, items)
		session.schedulePrefillSnapshots([]int{itemPos + itemLen/2})
		feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1])
		session.attachPrefillSnapshots()
		session.close()

		walkNodes(pc.root, func(n *trieNode) bool {
			if itemPos < n.endOffset && n.endOffset < itemPos+itemLen {
				t.Errorf("trie node ends at %d, inside the item's tokens [%d,%d)", n.endOffset, itemPos, itemPos+itemLen)
			}
			return true
		})
		if !nodeExistsAtOffset(pc.root, itemPos+itemLen) {
			t.Errorf("capture inside the item did not move to its end %d", itemPos+itemLen)
		}

		// A prompt ending inside the item matches the stored path through its
		// last token, which would put the resume point inside the item.
		short := inputs[:itemPos+itemLen-2]
		shortItems := []mediaItem{{pos: itemPos, length: len(short) - itemPos, fold: 1 << 31, item: &base.PreparedItem{}}}
		session = pc.begin(short, shortItems)
		if resumed := len(short) - len(session.remaining); resumed > itemPos {
			t.Errorf("resumed at %d, inside the item's tokens starting at %d", resumed, itemPos)
		}
		session.close()

		checkTrieInvariants(t, pc.root)
	})
}

// TestPrefillSnapshotsKeptOnCancel mirrors a prefill canceled after the caches
// captured interior snapshots but before the success-path attach ran. Closing
// the session attaches the crossed captures so a retry can resume from them,
// and drains the capture schedule; otherwise the next request's
// PrepareSnapshots would overwrite it without closing the captures, leaking
// them (caught by checkSnapshotLeaks in the env cleanup).
func TestPrefillSnapshotsKeptOnCancel(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc
		inputs := []int32{1, 2, 3, 4, 5}

		session := pc.begin(inputs, nil)
		session.schedulePrefillSnapshots([]int{3})
		// Cross offset 3 so the caches capture it, then close the session as a
		// canceled prefill would, before the success-path attach.
		feedAll(pc.caches, inputs[pc.minCacheOffset():3])
		session.close()

		// The crossed capture becomes a restore point for the retry.
		if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) {
			t.Errorf("no trie node at capture point %d after cancel", at)
		}

		// A second request re-prepares snapshots on the same caches: if the
		// pending ones were not drained, prepare() orphans them here.
		simulateRequest(t, pc, inputs, nil, 5)

		checkTrieInvariants(t, pc.root)
	})
}

func nodeExistsAtOffset(root *trieNode, offset int) bool {
	var found bool
	walkNodes(root, func(n *trieNode) bool {
		if n.endOffset == offset && n != root {
			found = true
		}
		return true
	})
	return found
}

func findUserNode(t *testing.T, pc *prefixCache) *trieNode {
	t.Helper()
	var found *trieNode
	walkNodes(pc.root, func(n *trieNode) bool {
		if n.user {
			found = n
		}
		return true
	})
	if found == nil {
		t.Fatal("no user-marked node found")
	}
	return found
}

func assertUserNodeExists(t *testing.T, pc *prefixCache, label string) {
	t.Helper()
	var exists bool
	walkNodes(pc.root, func(n *trieNode) bool {
		if n.user {
			exists = true
		}
		return true
	})
	if !exists {
		t.Fatalf("%s: no user-marked node found", label)
	}
}

// TestBranchSwitchRestoresCorrectState exercises switching back to an older
// branch after working on a different one, verifying that the restored cache
// state contains the correct token sequence for both rewindable and
// non-rewindable caches.
func TestBranchSwitchRestoresCorrectState(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc

		// Request A: [1,2,3,4,5] + generate [10,11]
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
		env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 10})

		// Request B: [1,2,3,6,7] — diverges at token 4
		simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{12, 13})
		env.assertAllTokens(t, "after B", []int32{1, 2, 3, 6, 7, 12})

		// Request C: switch back to A's branch [1,2,3,4,5,10,11,20]
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20}, nil)
		env.assertAllTokens(t, "after C (back to A)", []int32{1, 2, 3, 4, 5, 10, 11})

		checkTrieInvariants(t, pc.root)
	})
}

// TestLRUOnlyUpdatesUsedNodes verifies that intermediate nodes on the active
// path whose snapshots were not actually restored don't get their lastUsed
// refreshed, allowing them to age out and collapse.
func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
	forEachEnv(t, func(t *testing.T, env *testEnv) {
		pc := env.pc

		// Request A: creates path [1,2,3,4,5] + generate [10,11]
		simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})

		// Request B: diverges at token 4, creating a branch point at offset 3
		// with a split snapshot.
		simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20, 21})

		// Set all lastUsed to a known old time.
		oldTime := time.Now().Add(-1 * time.Hour)
		walkNodes(pc.root, func(n *trieNode) bool {
			n.lastUsed = oldTime
			return true
		})

		// Request C: continue on B's branch. This will match B's path
		// and extend it. The branch point's snapshot may be paged in
		// for some cache types but not others.
		beforeRequest := time.Now()
		inputsC := []int32{1, 2, 3, 6, 7, 20, 21, 30}
		resC := simulateRequest(t, pc, inputsC, nil)
		landing := len(inputsC) - len(resC.remaining)

		// The path must have enough depth to exercise intermediate nodes.
		if len(pc.activePath) < 3 {
			t.Fatalf("activePath too short to test intermediate nodes: got %d nodes", len(pc.activePath))
		}

		// The frontier (deepest node on the active path) must be updated.
		frontier := pc.activePath[len(pc.activePath)-1]
		if frontier.lastUsed.Before(beforeRequest) {
			t.Errorf("frontier lastUsed was not updated: got %v, want >= %v",
				frontier.lastUsed, beforeRequest)
		}

		// Only used nodes refresh — the frontier and the restore landing;
		// merely traversed nodes keep their age so they can still evict.
		for i, node := range pc.activePath[:len(pc.activePath)-1] {
			if node.endOffset == landing {
				continue
			}
			if !node.lastUsed.Before(beforeRequest) {
				t.Errorf("activePath[%d] (endOffset=%d) lastUsed was refreshed: got %v, want < %v",
					i, node.endOffset, node.lastUsed, beforeRequest)
			}
		}

		checkTrieInvariants(t, pc.root)
	})
}

// TestPagedOutBytesUpdatesOnMaterialize verifies that when a snapshot owned
// by a trie node materializes (allocates owned bytes from a previously lazy
// state), the trie's pagedOutBytes counter picks up the delta via the
// installed materialize hook.
func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
	pc := &prefixCache{}
	pc.ensureRoot()

	node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
	pc.root.children = append(pc.root.children, node)

	snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
	node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)

	if pc.pagedOutBytes != 0 {
		t.Fatalf("pagedOutBytes after install = %d, want 0 (lazy snapshot)", pc.pagedOutBytes)
	}

	const materialized = 1 << 20
	snap.materialize(materialized)

	if pc.pagedOutBytes != materialized {
		t.Fatalf("pagedOutBytes after materialize = %d, want %d", pc.pagedOutBytes, materialized)
	}
}

// TestSwapSnapshotsDetachesHook verifies that snapshots removed from a trie
// node via swapSnapshots no longer feed the trie's counter when they later
// materialize. Without detach, a Split/Merge that folds an old snapshot
// elsewhere would double-count its bytes if it copied out afterward.
func TestSwapSnapshotsDetachesHook(t *testing.T) {
	pc := &prefixCache{}
	pc.ensureRoot()

	node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
	pc.root.children = append(pc.root.children, node)

	snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
	node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)

	replacement := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
	old := node.swapSnapshots([]cache.Snapshot{replacement}, &pc.pagedOutBytes)
	if len(old) != 1 || old[0] != snap {
		t.Fatalf("swapSnapshots returned %v, want the original snap", old)
	}

	snap.materialize(1 << 20)
	if pc.pagedOutBytes != 0 {
		t.Fatalf("pagedOutBytes after detached materialize = %d, want 0", pc.pagedOutBytes)
	}

	replacement.materialize(2 << 20)
	if pc.pagedOutBytes != 2<<20 {
		t.Fatalf("pagedOutBytes after replacement materialize = %d, want %d", pc.pagedOutBytes, 2<<20)
	}
}

// Without this, every node reads as incomplete on a hybrid model and
// switchToPath re-pages the leaf every time.
func TestHasAllSnapshotsIgnoresStatelessLayers(t *testing.T) {
	tr := &snapshotTracker{}
	caches := []cache.Cache{
		&fakeRewindableCache{tracker: tr},
		nil,
		&fakeRecurrentCache{tracker: tr},
	}

	node := &trieNode{tokens: []trieKey{1, 2, 3}, endOffset: 3}
	if hasAllSnapshots(node, caches) {
		t.Fatal("hasAllSnapshots = true with no snapshots at all")
	}

	node.snapshots = []cache.Snapshot{&fakeSnapshot{from: 0, to: 3}, nil, &fakeSnapshot{from: 0, to: 3}}
	if !hasAllSnapshots(node, caches) {
		t.Fatal("hasAllSnapshots = false when every stateful layer has a snapshot")
	}

	node.snapshots = []cache.Snapshot{&fakeSnapshot{from: 0, to: 3}, nil, nil}
	if hasAllSnapshots(node, caches) {
		t.Fatal("hasAllSnapshots = true with a stateful layer's snapshot missing")
	}
}

// Merge is only reached when an interior node with one child is evicted, which
// the scenario tests never hit, so drive it directly.
func TestMergeWithChildSkipsStatelessLayers(t *testing.T) {
	tr := &snapshotTracker{}
	caches := []cache.Cache{
		&fakeRewindableCache{tracker: tr},
		nil,
		&fakeRecurrentCache{tracker: tr},
	}

	pc := &prefixCache{caches: caches}
	pc.ensureRoot()

	parent := &trieNode{parent: pc.root, tokens: []trieKey{1, 2}, endOffset: 2}
	child := &trieNode{parent: parent, tokens: []trieKey{3, 4}, endOffset: 4}
	parent.children = []*trieNode{child}
	pc.root.children = []*trieNode{parent}

	parent.setSnapshots([]cache.Snapshot{
		&fakeSnapshot{from: 0, to: 2}, nil, &fakeSnapshot{from: 0, to: 2},
	}, &pc.pagedOutBytes)
	child.setSnapshots([]cache.Snapshot{
		&fakeSnapshot{from: 2, to: 4}, nil, &fakeSnapshot{from: 2, to: 4},
	}, &pc.pagedOutBytes)

	mergeWithChild(parent, caches, &pc.pagedOutBytes)

	if got := len(parent.tokens); got != 4 {
		t.Fatalf("merged tokens = %d, want 4", got)
	}
	if got := parent.endOffset; got != 4 {
		t.Fatalf("merged endOffset = %d, want 4", got)
	}
	if parent.snapshots[1] != nil {
		t.Fatalf("stateless layer snapshot = %v, want nil", parent.snapshots[1])
	}
	for _, i := range []int{0, 2} {
		if parent.snapshots[i] == nil {
			t.Fatalf("stateful layer %d lost its merged snapshot", i)
		}
	}
}