Allocate node and symbol IDs in per-checker blocks

internal/api/session.go

@@ -41,7 +41,7 @@ type snapshotData struct {
 	snapshot *project.Snapshot
 	refCount int
 
-	// Symbol IDs come from ast.GetSymbolId, a global atomic counter, so the same
+	// Symbol IDs come from global atomic counters in ast, so the same
 	// *ast.Symbol pointer always has the same unique ID across all projects in the
 	// snapshot. Symbols are registered snapshot-wide to ensure identity semantics:
 	// querying the same symbol from two different projects returns the same handle.

internal/ast/ids.go

@@ -1,6 +1,100 @@
 package ast
 
+import (
+	"sync/atomic"
+
+	"github.com/microsoft/typescript-go/internal/core"
+)
+
 type (
 	NodeId   uint64
 	SymbolId uint64
 )
+
+// Atomic ids
+
+var (
+	nextNodeId   atomic.Uint64
+	nextSymbolId atomic.Uint64
+)
+
+func GetNodeId(node *Node) NodeId {
+	return NodeId(getId(&node.id, &nextNodeId))
+}
+
+func GetSymbolId(symbol *Symbol) SymbolId {
+	return SymbolId(getId(&symbol.id, &nextSymbolId))
+}
+
+func getId(id *atomic.Uint64, counter *atomic.Uint64) uint64 {
+	value := id.Load()
+	if value == 0 {
+		// Worst case, we burn a few ids if we have to CAS.
+		value = counter.Add(1)
+		if !id.CompareAndSwap(0, value) {
+			value = id.Load()
+		}
+	}
+	return value
+}
+
+// IDs are handed out in page-aligned blocks so that a link store page only ever holds IDs minted
+// by one allocator. Blocks grow from one page to maxIdBlockSize as an allocator keeps allocating.
+const (
+	minIdBlockSize = core.LinkStorePageSize
+	maxIdBlockSize = 16 * core.LinkStorePageSize
+)
+
+// IdAllocator assigns node and symbol IDs from contiguous blocks reserved from the global counters,
+// which keeps the IDs one client mints clustered in the same link store pages. An IdAllocator must
+// not be used concurrently.
+type IdAllocator struct {
+	nodeIds   idBlock
+	symbolIds idBlock
+}
+
+func (a *IdAllocator) GetNodeId(node *Node) NodeId {
+	if id := node.id.Load(); id != 0 {
+		return NodeId(id)
+	}
+	return NodeId(a.nodeIds.assign(&node.id, &nextNodeId))
+}
+
+func (a *IdAllocator) GetSymbolId(symbol *Symbol) SymbolId {
+	if id := symbol.id.Load(); id != 0 {
+		return SymbolId(id)
+	}
+	return SymbolId(a.symbolIds.assign(&symbol.id, &nextSymbolId))
+}
+
+type idBlock struct {
+	next uint64
+	end  uint64
+	size uint64
+}
+
+func (b *idBlock) assign(id *atomic.Uint64, counter *atomic.Uint64) uint64 {
+	if b.next == b.end {
+		b.reserve(counter)
+	}
+	value := b.next
+	if !id.CompareAndSwap(0, value) {
+		return id.Load()
+	}
+	b.next++
+	return value
+}
+
+// reserve claims the next page-aligned block of IDs from counter.
+func (b *idBlock) reserve(counter *atomic.Uint64) {
+	b.size = min(max(2*b.size, minIdBlockSize), maxIdBlockSize)
+	for {
+		last := counter.Load()
+		first := (last/core.LinkStorePageSize + 1) * core.LinkStorePageSize
+		if counter.CompareAndSwap(last, first+b.size-1) {
+			b.next = first
+			b.end = first + b.size
+			return
+		}
+	}
+}

internal/ast/ids_test.go

@@ -0,0 +1,29 @@
+package ast_test
+
+import (
+	"testing"
+
+	"github.com/microsoft/typescript-go/internal/ast"
+	"github.com/microsoft/typescript-go/internal/core"
+	"gotest.tools/v3/assert"
+)
+
+func TestIdAllocatorPagesAreExclusive(t *testing.T) {
+	t.Parallel()
+
+	var allocators [4]ast.IdAllocator
+	pages := make(map[ast.SymbolId]int)
+	for i := range 100000 {
+		a := i % len(allocators)
+		symbol := &ast.Symbol{}
+		id := allocators[a].GetSymbolId(symbol)
+		assert.Equal(t, ast.GetSymbolId(symbol), id)
+		page := id / core.LinkStorePageSize
+		owner, seen := pages[page]
+		if seen {
+			assert.Equal(t, owner, a, "page %d holds ids from two allocators", page)
+		} else {
+			pages[page] = a
+		}
+	}
+}

internal/ast/utilities.go

@@ -5,44 +5,12 @@ import (
 	"slices"
 	"strings"
 	"sync"
-	"sync/atomic"
 
 	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/debug"
 	"github.com/microsoft/typescript-go/internal/tspath"
 )
 
-// Atomic ids
-
-var (
-	nextNodeId   atomic.Uint64
-	nextSymbolId atomic.Uint64
-)
-
-func GetNodeId(node *Node) NodeId {
-	id := node.id.Load()
-	if id == 0 {
-		// Worst case, we burn a few ids if we have to CAS.
-		id = nextNodeId.Add(1)
-		if !node.id.CompareAndSwap(0, id) {
-			id = node.id.Load()
-		}
-	}
-	return NodeId(id)
-}
-
-func GetSymbolId(symbol *Symbol) SymbolId {
-	id := symbol.id.Load()
-	if id == 0 {
-		// Worst case, we burn a few ids if we have to CAS.
-		id = nextSymbolId.Add(1)
-		if !symbol.id.CompareAndSwap(0, id) {
-			id = symbol.id.Load()
-		}
-	}
-	return SymbolId(id)
-}
-
 func GetSymbolTable(data *SymbolTable) SymbolTable {
 	if *data == nil {
 		*data = make(SymbolTable)

internal/checker/checker.go

@@ -578,6 +578,7 @@ var nextCheckerID atomic.Uint32
 
 type Checker struct {
 	id                                          uint32
+	ids                                         ast.IdAllocator
 	program                                     Program
 	compilerOptions                             *core.CompilerOptions
 	files                                       []*ast.SourceFile
@@ -903,6 +904,8 @@ func NewChecker(program Program, tracer *Tracer) (*Checker, *sync.Mutex) {
 
 	c := &Checker{}
 	c.id = nextCheckerID.Add(1)
+	c.symbolNodeLinks.ids = &c.ids
+	c.valueSymbolLinks.ids = &c.ids
 	c.tracer = tracer
 	c.program = program
 	c.compilerOptions = program.Options()

internal/checker/links.go

@@ -9,29 +9,31 @@ import (
 // in the pages of the store which is suitable for values where sizeof(V) is small.
 type nodeLinkStore[V any] struct {
 	store core.PagedLinkStore[V]
+	ids   *ast.IdAllocator
 }
 
 func (s *nodeLinkStore[V]) Get(node *ast.Node) *V {
-	return s.store.Get(uint64(ast.GetNodeId(node)))
+	return s.store.Get(uint64(s.ids.GetNodeId(node)))
 }
 
 func (s *nodeLinkStore[V]) Has(node *ast.Node) bool {
-	return s.store.Has(uint64(ast.GetNodeId(node)))
+	return s.store.Has(uint64(s.ids.GetNodeId(node)))
 }
 
 func (s *nodeLinkStore[V]) TryGet(node *ast.Node) *V {
-	return s.store.TryGet(uint64(ast.GetNodeId(node)))
+	return s.store.TryGet(uint64(s.ids.GetNodeId(node)))
 }
 
 // symbolArenaLinkStore is a links store keyed by symbol references. Values are stored
 // indirectly in an arena which is suitable for values where sizeof(V) is larger.
 type symbolArenaLinkStore[V any] struct {
 	store core.PagedLinkStore[*V]
 	arena core.Arena[V]
+	ids   *ast.IdAllocator
 }
 
 func (s *symbolArenaLinkStore[V]) Get(symbol *ast.Symbol) *V {
-	link := s.store.Get(uint64(ast.GetSymbolId(symbol)))
+	link := s.store.Get(uint64(s.ids.GetSymbolId(symbol)))
 	if *link == nil {
 		*link = s.arena.New()
 	}
@@ -43,7 +45,7 @@ func (s *symbolArenaLinkStore[V]) Has(symbol *ast.Symbol) bool {
 }
 
 func (s *symbolArenaLinkStore[V]) TryGet(symbol *ast.Symbol) *V {
-	if link := s.store.TryGet(uint64(ast.GetSymbolId(symbol))); link != nil {
+	if link := s.store.TryGet(uint64(s.ids.GetSymbolId(symbol))); link != nil {
 		return *link
 	}
 	return nil

internal/core/linkstore.go

@@ -32,10 +32,11 @@ func (s *LinkStore[K, V]) TryGet(key K) *V {
 }
 
 const (
-	pageShift    = 8
-	pageSize     = 1 << pageShift
-	pageMask     = pageSize - 1
-	maxPageCount = 65536
+	pageShift         = 8
+	pageSize          = 1 << pageShift
+	pageMask          = pageSize - 1
+	maxPageCount      = 65536
+	LinkStorePageSize = pageSize
 )
 
 // Implements a sparse-array-like structure for storing elements keyed by dense uint64 keys. Elements are