fix(radls): refcount DocumentVersion to free tree-sitter C memory

Every parser.Parse allocates a *ts.Tree wrapping C-heap memory.
go-tree-sitter requires explicit Close() to release it (the
upstream README explicitly warns against finalizers due to cgo
hazards: Nodes hold references into the C tree but the Go GC
can't see them, so a finalizer-driven Close can fire while
another goroutine is walking the tree). The Phase 8b snapshot
model dropped the only call site that closed trees, so every
didChange now leaks one tree until the process exits.

Implement reference counting on DocumentVersion. The Document
holds one reference; each State.Snapshot / SnapshotByID call
bumps the count and returns to a caller who MUST call Release.
When the count reaches zero the tree is Close()d.

The acquire path uses CAS rather than plain Add so we can't
resurrect a snapshot whose tree has already been freed - the
'weak-to-strong reference upgrade' pattern. A small race window
exists between loading Document.snapshot and acquire(): the
writer can Release the old version we observed. acquire returns
false in that case and we retry; Document.snapshot has by then
been updated to point at a newer version.

Also fix the secondary leak in RadTree.Update (the standalone
checker path): the old tree was being abandoned without Close.
Now Update Closes the previous root.

The RadTree.Close becomes idempotent via sync.Once because the
underlying ts.Tree.Close is NOT idempotent (always calls
ts_tree_delete; second call double-frees). closeOnce lets
callers defer-Close confidently and also protects the
Document.Update -> Release path when Update's child tree replaces
a snapshot.

Tests: explicit refs-go-to-zero -> tree-nil verification,
late-acquire-on-released-snapshot must fail, all existing
snapshot/concurrency tests updated to Release. Race-clean under
-race across the full project.

Author: amterp

radls/analysis/differential_test.go

@@ -70,10 +70,12 @@ func TestDifferentialFinalStateMatchesFromScratch(t *testing.T) {
 				})
 			}
 			incSnap := incremental.Snapshot(uri)
+			defer incSnap.Release()
 
 			scratch := freshState()
 			scratch.AddDoc(uri, final)
 			scrSnap := scratch.Snapshot(uri)
+			defer scrSnap.Release()
 
 			assertEquivalent(t, incSnap, scrSnap)
 		})
@@ -106,8 +108,16 @@ func TestDifferentialEditCommutativityForIndependentDocs(t *testing.T) {
 	s2.UpdateDoc(uriB, []lsp.TextDocumentContentChangeEvent{{Text: textB}})
 	s2.UpdateDoc(uriA, []lsp.TextDocumentContentChangeEvent{{Text: textA}})
 
-	assertEquivalent(t, s1.Snapshot(uriA), s2.Snapshot(uriA))
-	assertEquivalent(t, s1.Snapshot(uriB), s2.Snapshot(uriB))
+	s1A := s1.Snapshot(uriA)
+	s2A := s2.Snapshot(uriA)
+	s1B := s1.Snapshot(uriB)
+	s2B := s2.Snapshot(uriB)
+	defer s1A.Release()
+	defer s2A.Release()
+	defer s1B.Release()
+	defer s2B.Release()
+	assertEquivalent(t, s1A, s2A)
+	assertEquivalent(t, s1B, s2B)
 }
 
 func freshState() *State {

radls/analysis/document.go

@@ -16,13 +16,25 @@ import (
 // concurrently without coordination: the data inside is guaranteed
 // not to mutate. The next didChange produces a NEW DocumentVersion
 // (built off the old) and atomically swaps the owning Document's
-// current pointer; old versions remain valid for any reader still
-// holding them, and are GC'd once unreferenced.
+// current pointer.
 //
 // This is the load-bearing piece of Phase 8: it lets LSP request
 // handlers (hover, goto-def, completion, etc.) grab a snapshot once
 // and reason about a frozen world for the duration of the request,
 // rather than racing against the next keystroke.
+//
+// Lifetime: the underlying tree-sitter tree owns C-heap memory that
+// the Go GC can't reclaim. We refcount accordingly: the Document
+// holds one reference; each call to State.Snapshot bumps the count
+// and returns the snapshot to the caller, who MUST call Release()
+// when done. When the count reaches zero, the tree is Close()d.
+//
+// Why refcounting and not a finalizer: tree-sitter Nodes carry
+// references into the C tree that the Go GC doesn't see. A
+// finalizer can fire on a tree while another goroutine is walking
+// nodes that point into it - the documented cgo hazard. Explicit
+// release gives us deterministic free at the moment we know no
+// reader is touching the tree.
 type DocumentVersion struct {
 	uri         string
 	fileID      FileID
@@ -32,17 +44,59 @@ type DocumentVersion struct {
 	ast         *rl.SourceFile
 	lineIndex   *LineIndex
 	diagnostics []lsp.Diagnostic
+
+	// refs starts at 1 - the Document that owns this version holds
+	// the initial reference. Each State.Snapshot caller bumps to one
+	// more; their matching Release drops it. When this hits zero the
+	// underlying tree is freed and any later acquire() returns false.
+	refs atomic.Int32
 }
 
-func (v *DocumentVersion) URI() string                  { return v.uri }
-func (v *DocumentVersion) FileID() FileID               { return v.fileID }
-func (v *DocumentVersion) Version() int64               { return v.version }
-func (v *DocumentVersion) Text() string                 { return v.text }
-func (v *DocumentVersion) Tree() *rts.RadTree           { return v.tree }
-func (v *DocumentVersion) AST() *rl.SourceFile          { return v.ast }
-func (v *DocumentVersion) LineIndex() *LineIndex        { return v.lineIndex }
+func (v *DocumentVersion) URI() string                   { return v.uri }
+func (v *DocumentVersion) FileID() FileID                { return v.fileID }
+func (v *DocumentVersion) Version() int64                { return v.version }
+func (v *DocumentVersion) Text() string                  { return v.text }
+func (v *DocumentVersion) Tree() *rts.RadTree            { return v.tree }
+func (v *DocumentVersion) AST() *rl.SourceFile           { return v.ast }
+func (v *DocumentVersion) LineIndex() *LineIndex         { return v.lineIndex }
 func (v *DocumentVersion) Diagnostics() []lsp.Diagnostic { return v.diagnostics }
 
+// acquire bumps the refcount if the snapshot is still live. Returns
+// false if the snapshot has already been released (refs == 0), in
+// which case the caller should retry the Document.Snapshot load -
+// the State has a newer version.
+//
+// Uses CAS rather than a plain Add so we never resurrect a snapshot
+// whose tree has already been Close()d. This is the standard
+// "weak-to-strong reference upgrade" pattern.
+func (v *DocumentVersion) acquire() bool {
+	for {
+		n := v.refs.Load()
+		if n == 0 {
+			return false
+		}
+		if v.refs.CompareAndSwap(n, n+1) {
+			return true
+		}
+	}
+}
+
+// Release drops one reference. When the count reaches zero the
+// underlying tree-sitter tree is freed. Each call to State.Snapshot
+// pairs with exactly one Release - callers typically `defer
+// snap.Release()` right after the nil-check.
+func (v *DocumentVersion) Release() {
+	if v == nil {
+		return
+	}
+	if v.refs.Add(-1) == 0 {
+		if v.tree != nil {
+			v.tree.Close()
+			v.tree = nil
+		}
+	}
+}
+
 // GetLine returns the source of the line at the given index, or "" if
 // out of range. Kept on DocumentVersion (not LineIndex) because callers
 // usually want both the text and the index together.
@@ -78,13 +132,19 @@ func (d *Document) Snapshot() *DocumentVersion {
 
 // Update runs `produce` under the writer lock to compute the next
 // version from the previous (nil on first open), then atomically swaps
-// it into place. Returns the new version.
+// it into place. The new version arrives with refs=1 (held by us);
+// after the store we Release the previous version, dropping
+// Document's reference to it. Any reader that had already Acquired
+// the old version keeps it alive via the refcount.
 func (d *Document) Update(produce func(prev *DocumentVersion) *DocumentVersion) *DocumentVersion {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 	prev := d.snapshot.Load()
 	next := produce(prev)
 	d.snapshot.Store(next)
+	if prev != nil {
+		prev.Release()
+	}
 	return next
 }
 
@@ -107,7 +167,7 @@ func buildVersion(
 	checker := check.NewCheckerWithTree(tree, parser, text, ast)
 	diags := runChecker(checker, lineIndex, encoding)
 
-	return &DocumentVersion{
+	v := &DocumentVersion{
 		uri:         uri,
 		fileID:      fileID,
 		version:     version,
@@ -117,6 +177,10 @@ func buildVersion(
 		lineIndex:   lineIndex,
 		diagnostics: diags,
 	}
+	// Owner's reference. Released by Document.Update when this
+	// version is replaced by a successor.
+	v.refs.Store(1)
+	return v
 }
 
 // runChecker is the boundary between check.Diagnostic (utf-8 byte

radls/analysis/document_test.go

@@ -10,7 +10,8 @@ import (
 // TestSnapshotStability exercises the central guarantee of Phase 8:
 // once a reader has a *DocumentVersion, subsequent writes don't change
 // what that pointer observes. The reader sees the world as it was
-// when it grabbed its snapshot.
+// when it grabbed its snapshot. The reader's Acquire keeps the
+// underlying tree alive across writer-driven version swaps.
 func TestSnapshotStability(t *testing.T) {
 	s := NewState()
 	s.SetEncoding(EncodingUTF16)
@@ -22,6 +23,7 @@ func TestSnapshotStability(t *testing.T) {
 	if first == nil {
 		t.Fatal("expected snapshot after AddDoc")
 	}
+	defer first.Release()
 	if got := first.Text(); got != "x = 1" {
 		t.Errorf("v1 text: got %q, want %q", got, "x = 1")
 	}
@@ -43,6 +45,7 @@ func TestSnapshotStability(t *testing.T) {
 	}
 
 	latest := s.Snapshot(uri)
+	defer latest.Release()
 	if latest == first {
 		t.Errorf("Snapshot() returned the same pointer after updates")
 	}
@@ -54,6 +57,52 @@ func TestSnapshotStability(t *testing.T) {
 	}
 }
 
+// TestSnapshotReleaseFreesTreeWhenLastRefDropped verifies the tree
+// is closed once the refcount reaches zero. Direct test of the
+// memory-leak fix: each Snapshot bumps refs, each Release drops one,
+// and when no more references exist the tree's C memory is freed.
+func TestSnapshotReleaseFreesTreeWhenLastRefDropped(t *testing.T) {
+	s := NewState()
+	s.SetEncoding(EncodingUTF16)
+	const uri = "file:///release.rad"
+	s.AddDoc(uri, "x = 1")
+
+	first := s.Snapshot(uri)
+	if first == nil {
+		t.Fatal("expected snapshot")
+	}
+	// At this point: refs = 2 (Document + caller).
+	if got := first.refs.Load(); got != 2 {
+		t.Errorf("after Snapshot: refs=%d, want 2", got)
+	}
+
+	// Update once. Document drops its reference to `first`, leaving
+	// just the caller's reference.
+	s.UpdateDoc(uri, []lsp.TextDocumentContentChangeEvent{{Text: "x = 2"}})
+	if got := first.refs.Load(); got != 1 {
+		t.Errorf("after Update: refs=%d, want 1", got)
+	}
+	if first.tree == nil {
+		t.Error("tree should still be alive while caller holds a ref")
+	}
+
+	// Caller releases. Refcount hits zero, tree is closed and set
+	// to nil.
+	first.Release()
+	if got := first.refs.Load(); got != 0 {
+		t.Errorf("after Release: refs=%d, want 0", got)
+	}
+	if first.tree != nil {
+		t.Error("tree should be nil after last Release")
+	}
+
+	// A late Acquire on a freed snapshot must fail. Without this the
+	// refcount could go negative and we'd never detect the bug.
+	if first.acquire() {
+		t.Error("acquire on released snapshot should fail")
+	}
+}
+
 // TestSnapshotConcurrentReaders runs many goroutines that all read
 // snapshots while writers churn. Doesn't assert on race conditions
 // directly (that's `go test -race`'s job), but does verify the data
@@ -102,6 +151,7 @@ func TestSnapshotConcurrentReaders(t *testing.T) {
 				txt := snap.Text()
 				_ = snap.Version()
 				_ = snap.LineIndex().LineCount()
+				snap.Release()
 				if len(txt) == 0 {
 					t.Errorf("empty text in snapshot")
 					return
@@ -112,7 +162,10 @@ func TestSnapshotConcurrentReaders(t *testing.T) {
 
 	// Let it run a bit, then stop.
 	for i := 0; i < 200; i++ {
-		_ = s.Snapshot(uri)
+		snap := s.Snapshot(uri)
+		if snap != nil {
+			snap.Release()
+		}
 	}
 	close(stop)
 	wg.Wait()

radls/analysis/fileid_test.go

@@ -42,6 +42,7 @@ func TestFileIDAssignment(t *testing.T) {
 	if snapA == nil {
 		t.Fatal("SnapshotByID(idA) returned nil")
 	}
+	defer snapA.Release()
 	if snapA.URI() != uriA {
 		t.Errorf("SnapshotByID resolved to wrong doc: got URI %q, want %q",
 			snapA.URI(), uriA)
@@ -76,8 +77,13 @@ func TestFileIDStableAcrossUpdates(t *testing.T) {
 	}
 
 	snap := s.SnapshotByID(id1)
-	if snap == nil || snap.Text() != "v2" {
-		t.Errorf("SnapshotByID after update: expected text v2, got %+v", snap)
+	if snap == nil {
+		t.Errorf("SnapshotByID after update: nil")
+	} else {
+		defer snap.Release()
+		if snap.Text() != "v2" {
+			t.Errorf("SnapshotByID after update: expected text v2, got %q", snap.Text())
+		}
 	}
 }
 

radls/analysis/state.go

@@ -65,16 +65,24 @@ func (s *State) FileIDFor(uri string) FileID {
 }
 
 // SnapshotByID returns the current snapshot for a FileID, or nil if
-// the id isn't known. Lets internal code that holds a FileID grab the
-// latest version without converting back to URI first.
+// the id isn't known. Same Acquire/Release contract as Snapshot:
+// caller MUST Release() what they get back.
 func (s *State) SnapshotByID(id FileID) *DocumentVersion {
 	s.mu.RLock()
 	doc, ok := s.idToDoc[id]
 	s.mu.RUnlock()
 	if !ok {
 		return nil
 	}
-	return doc.Snapshot()
+	for {
+		snap := doc.Snapshot()
+		if snap == nil {
+			return nil
+		}
+		if snap.acquire() {
+			return snap
+		}
+	}
 }
 
 // Encoding returns the LSP position encoding currently in use.
@@ -90,17 +98,34 @@ func (s *State) SetEncoding(enc PositionEncoding) {
 	s.encoding = enc
 }
 
-// Snapshot returns the current version of the named document, or nil
-// if the document isn't open. Lock-free on the version side; the docs
-// map lookup takes the RWMutex read lock briefly.
+// Snapshot returns the current version of the named document with
+// its refcount incremented. The caller MUST call Release() when
+// done (typically `defer snap.Release()` right after the nil-check).
+// Returns nil if the document isn't open.
+//
+// The acquire-after-load loop handles a small race window: between
+// loading the atomic pointer and bumping the refcount, the writer
+// could have Released the version we observed. In that case acquire
+// returns false and we retry; the Document.snapshot pointer has
+// already been updated to a newer version by the time we get here.
 func (s *State) Snapshot(uri string) *DocumentVersion {
 	s.mu.RLock()
 	doc, ok := s.docs[uri]
 	s.mu.RUnlock()
 	if !ok {
 		return nil
 	}
-	return doc.Snapshot()
+	for {
+		snap := doc.Snapshot()
+		if snap == nil {
+			return nil
+		}
+		if snap.acquire() {
+			return snap
+		}
+		// snap was released between Load and acquire; the writer
+		// must have stored a newer version. Retry.
+	}
 }
 
 // document returns the *Document handle, creating any missing entry is

radls/server/server.go

@@ -97,6 +97,7 @@ func (s *Server) handleDidOpen(_ context.Context, params json.RawMessage) (err e
 	// somehow the snapshot is gone we skip - notifyDiagnostics with an
 	// empty slice would clear any prior diagnostics, which is wrong.
 	if snap := s.s.Snapshot(uri); snap != nil {
+		defer snap.Release()
 		s.notifyDiagnostics(uri, snap.Diagnostics())
 	}
 	return
@@ -117,6 +118,7 @@ func (s *Server) handleDidChange(_ context.Context, params json.RawMessage) (err
 		// further keystrokes between trigger and fire will have
 		// produced newer versions, and we want the latest.
 		if snap := s.s.Snapshot(uri); snap != nil {
+			defer snap.Release()
 			s.notifyDiagnostics(uri, snap.Diagnostics())
 		}
 	})
@@ -132,6 +134,9 @@ func (s *Server) handleCompletion(_ context.Context, params json.RawMessage) (re
 	// Any subsequent didChange produces a new snapshot but this
 	// handler operates on the one it grabbed - frozen, race-free.
 	snap := s.s.Snapshot(completionParams.TextDocument.Uri)
+	if snap != nil {
+		defer snap.Release()
+	}
 	result, err = s.s.Complete(snap, completionParams.Position)
 	return
 }
@@ -142,6 +147,9 @@ func (s *Server) handleCodeAction(_ context.Context, params json.RawMessage) (re
 		return
 	}
 	snap := s.s.Snapshot(codeActionParams.TextDocument.Uri)
+	if snap != nil {
+		defer snap.Release()
+	}
 	result, err = s.s.CodeAction(snap, codeActionParams.Range)
 	return
 }