Dense link stores

tsc/internal/ast/utilities.go

@@ -15,8 +15,10 @@ import (
 // Atomic ids
 
 var (
-	nextNodeId   atomic.Uint64
-	nextSymbolId atomic.Uint64
+	nextNodeId        atomic.Uint64
+	nextSymbolId      atomic.Uint64
+	nextNodeBlockId   atomic.Uint64
+	nextSymbolBlockId atomic.Uint64
 )
 
 func GetNodeId(node *Node) NodeId {
@@ -43,6 +45,53 @@ func GetSymbolId(symbol *Symbol) SymbolId {
 	return SymbolId(id)
 }
 
+const (
+	BlockIdOffset = 0x1_0000_0000_0000
+	BlockIdSize   = core.LinkPageSize
+)
+
+type NodeIdGenerator struct {
+	nextId uint64
+	lastId uint64
+}
+
+func (a *NodeIdGenerator) GetNodeId(node *Node) NodeId {
+	id := node.id.Load()
+	if id == 0 {
+		if a.nextId == a.lastId {
+			a.nextId = nextNodeBlockId.Add(BlockIdSize) - BlockIdSize
+			a.lastId = a.nextId + BlockIdSize
+		}
+		id = a.nextId + BlockIdOffset
+		a.nextId++
+		if !node.id.CompareAndSwap(0, id) {
+			id = node.id.Load()
+		}
+	}
+	return NodeId(id)
+}
+
+type SymbolIdGenerator struct {
+	nextId uint64
+	lastId uint64
+}
+
+func (a *SymbolIdGenerator) GetSymbolId(symbol *Symbol) SymbolId {
+	id := symbol.id.Load()
+	if id == 0 {
+		if a.nextId == a.lastId {
+			a.nextId = nextSymbolBlockId.Add(BlockIdSize) - BlockIdSize
+			a.lastId = a.nextId + BlockIdSize
+		}
+		id = a.nextId + BlockIdOffset
+		a.nextId++
+		if !symbol.id.CompareAndSwap(0, id) {
+			id = symbol.id.Load()
+		}
+	}
+	return SymbolId(id)
+}
+
 func GetSymbolTable(data *SymbolTable) SymbolTable {
 	if *data == nil {
 		*data = make(SymbolTable)

tsc/internal/checker/checker.go

@@ -684,7 +684,7 @@ type Checker struct {
 	jsxElementLinks                             core.LinkStore[*ast.Node, JsxElementLinks]
 	computedNameLinks                           core.LinkStore[*ast.Node, ComputedNameNodeLinks]
 	symbolReferenceLinks                        core.LinkStore[*ast.Symbol, SymbolReferenceLinks]
-	valueSymbolLinks                            symbolArenaLinkStore[ValueSymbolLinks]
+	valueSymbolLinks                            symbolLinkStore[ValueSymbolLinks]
 	mappedSymbolLinks                           core.LinkStore[*ast.Symbol, MappedSymbolLinks]
 	deferredSymbolLinks                         core.LinkStore[*ast.Symbol, DeferredSymbolLinks]
 	aliasSymbolLinks                            core.LinkStore[*ast.Symbol, AliasSymbolLinks]

tsc/internal/checker/links.go

@@ -5,46 +5,62 @@ import (
 	"github.com/microsoft/TypeScript/tsc/internal/core"
 )
 
-// nodeLinkStore is a links store keyed by node references. Values are stored directly
-// in the pages of the store which is suitable for values where sizeof(V) is small.
+// When possible, nodeLinkStore and symbolLinkStore store Node and Symbol links in an efficient and densely
+// packed paged array store. When a Node or Symbol has not yet been assigned an ID, the store provides one
+// from a generator that produces sequences of IDs in a reserved range of the ID space. The generator grabs
+// chunks of 256 IDs from a central atomic counter, ensuring that each block of 256 IDs is consecutive and
+// causing Node or Symbol links to be densely packed within a single page.
+
+const maxPageLinkCount = 0x100_0000 // 16M
+
 type nodeLinkStore[V any] struct {
-	store core.PagedLinkStore[V]
+	gen   ast.NodeIdGenerator
+	links core.LinkStore[uint64, V]
+	pages core.PagedLinkStore[V]
 }
 
 func (s *nodeLinkStore[V]) Get(node *ast.Node) *V {
-	return s.store.Get(uint64(ast.GetNodeId(node)))
+	id := uint64(s.gen.GetNodeId(node))
+	if id >= ast.BlockIdOffset && id < ast.BlockIdOffset+maxPageLinkCount {
+		return s.pages.Get(id - ast.BlockIdOffset)
+	}
+	return s.links.Get(id)
 }
 
 func (s *nodeLinkStore[V]) Has(node *ast.Node) bool {
-	return s.store.Has(uint64(ast.GetNodeId(node)))
+	return s.TryGet(node) != nil
 }
 
 func (s *nodeLinkStore[V]) TryGet(node *ast.Node) *V {
-	return s.store.TryGet(uint64(ast.GetNodeId(node)))
+	id := uint64(s.gen.GetNodeId(node))
+	if id >= ast.BlockIdOffset && id < ast.BlockIdOffset+maxPageLinkCount {
+		return s.pages.TryGet(id - ast.BlockIdOffset)
+	}
+	return s.links.TryGet(id)
 }
 
-// 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]
+type symbolLinkStore[V any] struct {
+	gen   ast.SymbolIdGenerator
+	links core.LinkStore[uint64, V]
+	pages core.PagedLinkStore[V]
 }
 
-func (s *symbolArenaLinkStore[V]) Get(symbol *ast.Symbol) *V {
-	link := s.store.Get(uint64(ast.GetSymbolId(symbol)))
-	if *link == nil {
-		*link = s.arena.New()
+func (s *symbolLinkStore[V]) Get(symbol *ast.Symbol) *V {
+	id := uint64(s.gen.GetSymbolId(symbol))
+	if id >= ast.BlockIdOffset && id < ast.BlockIdOffset+maxPageLinkCount {
+		return s.pages.Get(id - ast.BlockIdOffset)
 	}
-	return *link
+	return s.links.Get(id)
 }
 
-func (s *symbolArenaLinkStore[V]) Has(symbol *ast.Symbol) bool {
+func (s *symbolLinkStore[V]) Has(symbol *ast.Symbol) bool {
 	return s.TryGet(symbol) != nil
 }
 
-func (s *symbolArenaLinkStore[V]) TryGet(symbol *ast.Symbol) *V {
-	if link := s.store.TryGet(uint64(ast.GetSymbolId(symbol))); link != nil {
-		return *link
+func (s *symbolLinkStore[V]) TryGet(symbol *ast.Symbol) *V {
+	id := uint64(s.gen.GetSymbolId(symbol))
+	if id >= ast.BlockIdOffset && id < ast.BlockIdOffset+maxPageLinkCount {
+		return s.pages.TryGet(id - ast.BlockIdOffset)
 	}
-	return nil
+	return s.links.TryGet(id)
 }

tsc/internal/core/linkstore.go

@@ -1,6 +1,8 @@
 package core
 
-import "slices"
+import (
+	"slices"
+)
 
 // Links store
 
@@ -32,62 +34,49 @@ func (s *LinkStore[K, V]) TryGet(key K) *V {
 }
 
 const (
-	pageShift    = 8
-	pageSize     = 1 << pageShift
-	pageMask     = pageSize - 1
-	maxPageCount = 65536
+	LinkPageShift = 8
+	LinkPageSize  = 1 << LinkPageShift
+	LinkPageMask  = LinkPageSize - 1
 )
 
-// Implements a sparse-array-like structure for storing elements keyed by dense uint64 keys. Elements are
-// stored in fixed-size pages of 256 entries and an index of pages is maintained in an array for lower valued
-// page indices and a map for higher valued page indices.
+// PagedLinkStore implements a sparse-array-like structure for storing elements keyed by dense uint64 keys.
+// Elements are allocated in an area, element references are stored in fixed-size pages of 256 entries, and
+// an index of pages is maintained in a growable list.
+
 type PagedLinkStore[V any] struct {
-	pageMap  map[uint64]*[pageSize]V // Page map for page indices above maxPageCount
-	pageList []*[pageSize]V          // Page table for page indices below maxPageCount
+	pages []*[LinkPageSize]*V
+	arena Arena[V]
 }
 
 func (s *PagedLinkStore[V]) Get(key uint64) *V {
-	var page *[pageSize]V
-	pageIndex := key >> pageShift
-	if pageIndex < maxPageCount {
-		if int(pageIndex) >= len(s.pageList) {
-			// Grow the length of the list to pageIndex+1
-			s.pageList = slices.Grow(s.pageList, int(pageIndex)-len(s.pageList)+1)[:pageIndex+1]
-		}
-		page = s.pageList[pageIndex]
-		if page == nil {
-			page = new([pageSize]V)
-			s.pageList[pageIndex] = page
-		}
-	} else {
-		page = s.pageMap[pageIndex]
-		if page == nil {
-			page = new([pageSize]V)
-			if s.pageMap == nil {
-				s.pageMap = make(map[uint64]*[pageSize]V)
-			}
-			s.pageMap[pageIndex] = page
-		}
+	pageIndex := key >> LinkPageShift
+	if int(pageIndex) >= len(s.pages) {
+		// Grow the length of the list to pageIndex+1
+		s.pages = slices.Grow(s.pages, int(pageIndex)-len(s.pages)+1)[:pageIndex+1]
+	}
+	page := s.pages[pageIndex]
+	if page == nil {
+		page = new([LinkPageSize]*V)
+		s.pages[pageIndex] = page
 	}
-	return &page[key&pageMask]
+	link := page[key&LinkPageMask]
+	if link == nil {
+		link = s.arena.New()
+		page[key&LinkPageMask] = link
+	}
+	return link
 }
 
 func (s *PagedLinkStore[V]) Has(key uint64) bool {
 	return s.TryGet(key) != nil
 }
 
 func (s *PagedLinkStore[V]) TryGet(key uint64) *V {
-	var page *[pageSize]V
-	pageIndex := key >> pageShift
-	if pageIndex < maxPageCount {
-		if int(pageIndex) < len(s.pageList) {
-			page = s.pageList[pageIndex]
+	pageIndex := key >> LinkPageShift
+	if int(pageIndex) < len(s.pages) {
+		if page := s.pages[pageIndex]; page != nil {
+			return page[key&LinkPageMask]
 		}
-	} else {
-		page = s.pageMap[pageIndex]
-	}
-	if page != nil {
-		return &page[key&pageMask]
 	}
 	return nil
 }