fix(index): keep node line index sub-linear for single-line JSON documents

index/map_index_nodes.go

@@ -4,6 +4,9 @@
 package index
 
 import (
+	"cmp"
+	"slices"
+
 	"go.yaml.in/yaml/v4"
 )
 
@@ -41,8 +44,10 @@ type NodeOrigin struct {
 	Index *SpecIndex `json:"-" yaml:"-"`
 }
 
-// nodeLineEntry is a single (column, node) pair on one line of the spec. Lines hold very
-// few nodes, so a small slice scanned linearly is far cheaper than a per-line map.
+// nodeLineEntry is a single (column, node) pair on one line of the spec. Most lines hold
+// very few nodes, but a minified JSON document puts every node on one line, so per-line
+// work must stay sub-linear: entries are appended during the build and each line is then
+// sorted by column exactly once (see sortNodeLines), making lookups a binary search.
 type nodeLineEntry struct {
 	column int32
 	node   *yaml.Node
@@ -66,17 +71,21 @@ func (index *SpecIndex) awaitNodeMap() {
 	}
 }
 
-// lookupNodeLines returns the node stored at line/column, or nil if absent.
+// lookupNodeLines returns the node stored at line/column, or nil if absent. Lines must
+// be sorted by column (see sortNodeLines).
 func lookupNodeLines(lines [][]nodeLineEntry, line, column int) *yaml.Node {
 	if line < 0 || line >= len(lines) {
 		return nil
 	}
-	for _, e := range lines[line] {
-		if int(e.column) == column {
-			return e.node
-		}
+	i, found := slices.BinarySearchFunc(lines[line], int32(column), compareEntryColumn)
+	if !found {
+		return nil
 	}
-	return nil
+	return lines[line][i].node
+}
+
+func compareEntryColumn(e nodeLineEntry, column int32) int {
+	return cmp.Compare(e.column, column)
 }
 
 // MapNodes maps all nodes in the document by line and column. The index is built into a
@@ -90,6 +99,7 @@ func (index *SpecIndex) MapNodes(rootNode *yaml.Node) {
 	// lines are 1-based; +1 so line NumLines is directly addressable.
 	lines := make([][]nodeLineEntry, sizeHint+1)
 	lines = mapNodesRecursive(rootNode, lines)
+	sortNodeLines(lines)
 	index.nodeMapLock.Lock()
 	index.nodeLines = lines
 	index.nodeMapLock.Unlock()
@@ -107,9 +117,9 @@ func mapNodesRecursive(node *yaml.Node, lines [][]nodeLineEntry) [][]nodeLineEnt
 	return addNodeLineEntry(lines, node)
 }
 
-// addNodeLineEntry records node at its line/column, preserving the previous map
-// semantics: a later write to the same line/column replaces the earlier one
-// (parents are written after their children, so parents win collisions).
+// addNodeLineEntry appends node at its line. It never scans the line, so building the
+// index stays linear even when the whole document sits on a single line; duplicate
+// columns are collapsed later by sortNodeLines.
 func addNodeLineEntry(lines [][]nodeLineEntry, node *yaml.Node) [][]nodeLineEntry {
 	line := node.Line
 	if line < 0 {
@@ -124,13 +134,29 @@ func addNodeLineEntry(lines [][]nodeLineEntry, node *yaml.Node) [][]nodeLineEntr
 		copy(expanded, lines)
 		lines = expanded
 	}
-	entries := lines[line]
-	for i := range entries {
-		if int(entries[i].column) == node.Column {
-			entries[i].node = node
-			return lines
+	lines[line] = append(lines[line], nodeLineEntry{column: int32(node.Column), node: node})
+	return lines
+}
+
+// sortNodeLines orders every line by column and collapses duplicate columns, keeping the
+// entry written last. This preserves the original map semantics: mapNodesRecursive writes
+// parents after their children, so a parent sharing a position with its first child wins.
+func sortNodeLines(lines [][]nodeLineEntry) {
+	for i, entries := range lines {
+		if len(entries) < 2 {
+			continue
 		}
+		slices.SortStableFunc(entries, func(a, b nodeLineEntry) int {
+			return cmp.Compare(a.column, b.column)
+		})
+		out := entries[:1]
+		for _, e := range entries[1:] {
+			if e.column == out[len(out)-1].column {
+				out[len(out)-1] = e
+				continue
+			}
+			out = append(out, e)
+		}
+		lines[i] = out
 	}
-	lines[line] = append(entries, nodeLineEntry{column: int32(node.Column), node: node})
-	return lines
 }

index/map_index_nodes_test.go

@@ -4,8 +4,10 @@
 package index
 
 import (
+	"fmt"
 	"os"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -117,11 +119,85 @@ func TestSpecIndex_MapNodes_OverwriteSemantics(t *testing.T) {
 	var lines [][]nodeLineEntry
 	lines = addNodeLineEntry(lines, first)
 	lines = addNodeLineEntry(lines, second)
+	sortNodeLines(lines)
 
 	assert.Same(t, second, lookupNodeLines(lines, 4, 2))
 	assert.Len(t, lines[4], 1)
 }
 
+func TestSpecIndex_SortNodeLines_OrdersAndDedupes(t *testing.T) {
+	// entries arrive in document (DFS) order, not column order. after sorting, lookups
+	// must work for every column and the last write for a column must win.
+	nodes := []*yaml.Node{
+		{Line: 1, Column: 30, Value: "c"},
+		{Line: 1, Column: 10, Value: "a"},
+		{Line: 1, Column: 20, Value: "b-child"},
+		{Line: 1, Column: 20, Value: "b-parent"},
+		{Line: 1, Column: 5, Value: "first"},
+	}
+	var lines [][]nodeLineEntry
+	for _, n := range nodes {
+		lines = addNodeLineEntry(lines, n)
+	}
+	sortNodeLines(lines)
+
+	assert.Len(t, lines[1], 4)
+	assert.Equal(t, "first", lines[1][0].node.Value, "first entry on a line is the leftmost node")
+	assert.Same(t, nodes[4], lookupNodeLines(lines, 1, 5))
+	assert.Same(t, nodes[1], lookupNodeLines(lines, 1, 10))
+	assert.Same(t, nodes[3], lookupNodeLines(lines, 1, 20))
+	assert.Same(t, nodes[0], lookupNodeLines(lines, 1, 30))
+	assert.Nil(t, lookupNodeLines(lines, 1, 25))
+}
+
+// singleLineJSONSpec builds a minified OpenAPI document with roughly n schema nodes,
+// all on line 1, mirroring specs published as compact JSON (e.g. very large public APIs).
+func singleLineJSONSpec(n int) []byte {
+	var b strings.Builder
+	b.WriteString(`{"openapi":"3.1.0","info":{"title":"dense","version":"1"},"paths":{},"components":{"schemas":{`)
+	for i := 0; i < n; i++ {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		fmt.Fprintf(&b, `"s%d":{"type":"object","properties":{"id":{"type":"integer"},"ref":{"$ref":"#/components/schemas/s%d"}}}`, i, (i+1)%n)
+	}
+	b.WriteString(`}}}`)
+	return []byte(b.String())
+}
+
+func TestSpecIndex_MapNodes_SingleLineDocument(t *testing.T) {
+	spec := singleLineJSONSpec(2000)
+	var rootNode yaml.Node
+	assert.NoError(t, yaml.Unmarshal(spec, &rootNode))
+
+	index := NewSpecIndexWithConfig(&rootNode, CreateOpenAPIIndexConfig())
+	<-index.nodeMapCompleted
+
+	// every node in the document lives on line 1 and must be retrievable by position.
+	total := 0
+	var walk func(n *yaml.Node)
+	walk = func(n *yaml.Node) {
+		for _, c := range n.Content {
+			walk(c)
+		}
+		if n.Kind == yaml.DocumentNode {
+			return
+		}
+		total++
+		assert.Equal(t, 1, n.Line)
+		found, ok := index.GetNode(n.Line, n.Column)
+		assert.True(t, ok)
+		assert.NotNil(t, found)
+	}
+	walk(&rootNode)
+	assert.Greater(t, total, 2000*8)
+
+	// the line index and the legacy map must agree.
+	legacy := index.GetNodeMap()
+	assert.Len(t, legacy, 1)
+	assert.Len(t, legacy[1], len(index.nodeLines[1]))
+}
+
 func TestSpecIndex_GetNodeMap_LegacyMaterialization(t *testing.T) {
 	petstore, _ := os.ReadFile("../test_specs/petstorev3.json")
 	var rootNode yaml.Node
@@ -166,6 +242,26 @@ func TestSpecIndex_GetNodeMap_AfterRelease(t *testing.T) {
 	assert.Nil(t, node)
 }
 
+// BenchmarkSpecIndex_MapNodes_SingleLine guards against per-line work becoming linear
+// again: with 20k schemas (~200k nodes on one line) a linear scan per insert or lookup
+// takes minutes, the sorted index takes milliseconds.
+func BenchmarkSpecIndex_MapNodes_SingleLine(b *testing.B) {
+	spec := singleLineJSONSpec(20000)
+	var rootNode yaml.Node
+	_ = yaml.Unmarshal(spec, &rootNode)
+	probe := rootNode.Content[0].Content[1] // "info" key node
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		index := NewSpecIndexWithConfig(&rootNode, CreateOpenAPIIndexConfig())
+		<-index.nodeMapCompleted
+		found, ok := index.GetNode(probe.Line, probe.Column)
+		if !ok || found != probe {
+			b.Fatal("probe node not found")
+		}
+	}
+}
+
 func BenchmarkSpecIndex_MapNodes(b *testing.B) {
 	petstore, _ := os.ReadFile("../test_specs/petstorev3.json")
 	var rootNode yaml.Node