Intern the Nodes that carry origin data

decode.go

@@ -377,6 +377,10 @@ type decoder struct {
 	aliasDepth        int
 	disableTimestamps bool
 
+	// origins interns the Nodes that encode origin data, which repeat heavily
+	// within a document. Created only when origin tracking is on.
+	origins *originCache
+
 	mergedFields map[interface{}]bool
 }
 
@@ -592,7 +596,7 @@ func (d *decoder) document(n *Node, out reflect.Value) (good bool) {
 					Line:   firstKey.Line,
 					Column: firstKey.Column,
 				}
-				addOriginInMap(syntheticKey, root, d.file)
+				addOriginInMap(syntheticKey, root, d.file, d.origins)
 			}
 		}
 		d.unmarshal(n.Content[0], out)
@@ -822,7 +826,7 @@ func (d *decoder) sequence(n *Node, out reflect.Value) (good bool) {
 	for i := 0; i < l; i++ {
 		e := reflect.New(et).Elem()
 		if d.origin && d.aliasDepth == 0 {
-			addOriginInSeq(n.Content[i], d.file)
+			addOriginInSeq(n.Content[i], d.file, d.origins)
 		}
 		if ok := d.unmarshal(n.Content[i], e); ok {
 			out.Index(j).Set(e)
@@ -939,7 +943,7 @@ func (d *decoder) mapping(n *Node, out reflect.Value) (good bool) {
 			e := reflect.New(et).Elem()
 
 			if d.origin && d.aliasDepth == 0 {
-				addOriginInMap(n.Content[i], n.Content[i+1], d.file)
+				addOriginInMap(n.Content[i], n.Content[i+1], d.file, d.origins)
 			}
 			if d.unmarshal(n.Content[i+1], e) || n.Content[i+1].ShortTag() == nullTag && (mapIsNew || !out.MapIndex(k).IsValid()) {
 				out.SetMapIndex(k, e)

origin.go

@@ -1,9 +1,68 @@
 package yaml
 
-import "fmt"
+import "strconv"
 
 const originTag = "__origin__"
 
+// Origin data is encoded as Nodes so it survives the node -> JSON ->
+// UnmarshalJSON conversion the decoder performs on the way to the caller's
+// type. That makes every recorded line, column, count and field name its own
+// heap-allocated Node carrying its own string, and on a large document those
+// dominate the decode: a profile of a 23 MB spec put intNode alone at 486 MB
+// of 3.07 GB allocated.
+//
+// These nodes are immutable leaf scalars -- nothing appends to their Content
+// or rewrites their Value -- so equal values can share one Node. Two caches
+// cover the two ways a spec repeats itself.
+//
+// Neither cache needs a size bound. The strings interned are mapping keys and
+// scalar sequence items (not values like descriptions), so the distinct set is
+// small by construction, and in the worst case the cache costs one pointer per
+// Node that would have been allocated regardless.
+
+// maxCachedInt covers essentially every column, line delta and count. Only an
+// absolute line number in a large document exceeds it, and there is one of
+// those per mapping against many small ones.
+const maxCachedInt = 1024
+
+var smallIntNodes = func() [maxCachedInt]*Node {
+	var nodes [maxCachedInt]*Node
+	for i := range nodes {
+		nodes[i] = &Node{Kind: ScalarNode, Tag: "!!int", Value: strconv.Itoa(i)}
+	}
+	return nodes
+}()
+
+// originKeyNode is the "__origin__" key appended to every mapping that gets
+// origin data. Its value is a constant, so one shared Node serves the whole
+// process. Line stays 0, which is what isOrigin tests for.
+var originKeyNode = &Node{Kind: ScalarNode, Tag: "!!str", Value: originTag}
+
+// originCache interns the nodes for one decode. It is per-decode rather than
+// global so it needs no locking and is reclaimed with the decoder.
+//
+// It keys on the string itself and holds no notion of "the" file, because the
+// file recorded in an origin is not constant: a $ref carries the decode into
+// another document, and the origin has to name the file the element actually
+// came from. Interning by value is correct either way -- a second file is
+// simply a second entry.
+type originCache struct {
+	strs map[string]*Node
+}
+
+func newOriginCache() *originCache {
+	return &originCache{strs: make(map[string]*Node)}
+}
+
+func (c *originCache) str(v string) *Node {
+	if n, ok := c.strs[v]; ok {
+		return n
+	}
+	n := &Node{Kind: ScalarNode, Tag: "!!str", Value: v}
+	c.strs[v] = n
+	return n
+}
+
 func isScalar(n *Node) bool {
 	return n.Kind == ScalarNode
 }
@@ -16,19 +75,19 @@ func isMapping(n *Node) bool {
 	return n.Kind == MappingNode
 }
 
-func addOriginInSeq(n *Node, file string) *Node {
+func addOriginInSeq(n *Node, file string, c *originCache) *Node {
 	if !isMapping(n) || len(n.Content) == 0 {
 		return n
 	}
 	// in case of a sequence, we use the first element as the key
-	return addOrigin(n.Content[0], n, file)
+	return addOrigin(n.Content[0], n, file, c)
 }
 
-func addOriginInMap(key, n *Node, file string) *Node {
+func addOriginInMap(key, n *Node, file string, c *originCache) *Node {
 	if !isMapping(n) {
 		return n
 	}
-	return addOrigin(key, n, file)
+	return addOrigin(key, n, file, c)
 }
 
 // addOrigin injects a compact __origin__ sequence into the mapping node n.
@@ -46,24 +105,25 @@ func addOriginInMap(key, n *Node, file string) *Node {
 //     key_line and absolute column of the position just past its last content.
 //     Appended last so a consumer that stops after the sequences section
 //     simply ignores it (backward compatible).
-func addOrigin(key, n *Node, file string) *Node {
+func addOrigin(key, n *Node, file string, c *originCache) *Node {
 	if isOrigin(key) {
 		return n
 	}
 
-	seq := buildOriginSeq(key, n, file)
+	seq := buildOriginSeq(key, n, file, c)
 	n.Content = append(n.Content,
-		&Node{Kind: ScalarNode, Tag: "!!str", Value: originTag}, // Line==0 → isOrigin
+		originKeyNode,
+		// Unlike the key, this one has to be fresh: it owns seq.
 		&Node{Kind: SequenceNode, Tag: "!!seq", Content: seq},
 	)
 	return n
 }
 
-func buildOriginSeq(key, n *Node, file string) []*Node {
+func buildOriginSeq(key, n *Node, file string, c *originCache) []*Node {
 	// Header: file, key_name, key_line, key_col
 	nodes := []*Node{
-		strNode(file),
-		strNode(key.Value),
+		c.str(file),
+		c.str(key.Value),
 		intNode(key.Line),
 		intNode(key.Column),
 	}
@@ -83,7 +143,7 @@ func buildOriginSeq(key, n *Node, file string) []*Node {
 		// Record the location of this field's key.
 		nf++
 		fieldNodes = append(fieldNodes,
-			strNode(k.Value),
+			c.str(k.Value),
 			intNode(k.Line-key.Line),
 			intNode(k.Column),
 		)
@@ -94,15 +154,15 @@ func buildOriginSeq(key, n *Node, file string) []*Node {
 			for _, item := range v.Content {
 				if item.Kind == ScalarNode {
 					itemNodes = append(itemNodes,
-						strNode(item.Value),
+						c.str(item.Value),
 						intNode(item.Line-key.Line),
 						intNode(item.Column),
 					)
 				}
 			}
 			if len(itemNodes) > 0 {
 				ns++
-				seqNodes = append(seqNodes, strNode(k.Value), intNode(len(itemNodes)/3))
+				seqNodes = append(seqNodes, c.str(k.Value), intNode(len(itemNodes)/3))
 				seqNodes = append(seqNodes, itemNodes...)
 			}
 		}
@@ -131,10 +191,9 @@ func isOrigin(key *Node) bool {
 	return key.Line == 0
 }
 
-func strNode(v string) *Node {
-	return &Node{Kind: ScalarNode, Tag: "!!str", Value: v}
-}
-
 func intNode(v int) *Node {
-	return &Node{Kind: ScalarNode, Tag: "!!int", Value: fmt.Sprintf("%d", v)}
+	if 0 <= v && v < maxCachedInt {
+		return smallIntNodes[v]
+	}
+	return &Node{Kind: ScalarNode, Tag: "!!int", Value: strconv.Itoa(v)}
 }

origin_intern_test.go

@@ -0,0 +1,104 @@
+package yaml_test
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"testing"
+
+	yaml "github.com/oasdiff/yaml3"
+	. "gopkg.in/check.v1"
+)
+
+const internDoc = "root:\n    hello: world\n    object:\n        foo: bar\n"
+
+func decodeWithOrigin(c *C, file string) string {
+	dec := yaml.NewDecoder(bytes.NewBufferString(internDoc))
+	dec.Origin(true, file)
+	var out any
+	c.Assert(dec.Decode(&out), IsNil)
+	b, err := yaml.Marshal(out)
+	c.Assert(err, IsNil)
+	return string(b)
+}
+
+// Origin nodes are interned, so a value occurring many times is one shared
+// Node. The file name is the value most at risk from that: it is recorded on
+// every mapping, and it is not constant across a document set, because a $ref
+// carries the origin into another file which is decoded separately. A cache
+// keyed on anything other than the string itself would smear one file's name
+// over another's origins.
+func (s *S) TestOrigin_InterningKeepsFilesDistinct(c *C) {
+	base := decodeWithOrigin(c, "base.yaml")
+	revision := decodeWithOrigin(c, "revision.json")
+
+	c.Assert(strings.Contains(base, "base.yaml"), Equals, true)
+	c.Assert(strings.Contains(base, "revision.json"), Equals, false)
+	c.Assert(strings.Contains(revision, "revision.json"), Equals, true)
+	c.Assert(strings.Contains(revision, "base.yaml"), Equals, false)
+}
+
+// Shared nodes are only safe while nothing mutates them. Decoding the same
+// input twice must reproduce it byte for byte, so a decode that wrote through
+// to a cached Node fails here rather than corrupting an unrelated document.
+func (s *S) TestOrigin_InternedNodesAreNotMutated(c *C) {
+	first := decodeWithOrigin(c, "base.yaml")
+	decodeWithOrigin(c, "other.yaml")
+	c.Assert(decodeWithOrigin(c, "base.yaml"), Equals, first)
+}
+
+// Line numbers past the small-int cache are allocated per use. Pin that they
+// are still recorded correctly, since the cached and uncached paths build the
+// Node differently.
+func (s *S) TestOrigin_LineBeyondTheSmallIntCache(c *C) {
+	var sb strings.Builder
+	sb.WriteString("root:\n")
+	for i := 0; i < 2000; i++ {
+		fmt.Fprintf(&sb, "    filler%d: x\n", i)
+	}
+	sb.WriteString("    tail:\n        leaf: value\n")
+
+	dec := yaml.NewDecoder(bytes.NewBufferString(sb.String()))
+	dec.Origin(true, "big.yaml")
+	var out any
+	c.Assert(dec.Decode(&out), IsNil)
+
+	root, ok := out.(map[string]any)["root"].(map[string]any)
+	c.Assert(ok, Equals, true)
+	tail, ok := root["tail"].(map[string]any)
+	c.Assert(ok, Equals, true)
+	origin, ok := tail["__origin__"].([]any)
+	c.Assert(ok, Equals, true)
+
+	// Header layout is [file, key_name, key_line, key_col, ...]. The key sits
+	// past filler0..filler1999, so its line exceeds maxCachedInt.
+	c.Assert(origin[0], Equals, "big.yaml")
+	c.Assert(origin[1], Equals, "tail")
+	c.Assert(toAnyInt(origin[2]), Equals, 2002)
+}
+
+func benchDoc(endpoints int) string {
+	var sb strings.Builder
+	sb.WriteString("openapi: 3.0.3\npaths:\n")
+	for i := 0; i < endpoints; i++ {
+		fmt.Fprintf(&sb, "    /resource/%d:\n        get:\n            operationId: get%d\n"+
+			"            responses:\n                \"200\":\n                    description: ok\n", i, i)
+	}
+	return sb.String()
+}
+
+// The origin nodes dominate allocation on a large document, which is what the
+// interning targets. Run with -benchmem; B/op is the number that moved.
+func BenchmarkOriginDecode(b *testing.B) {
+	doc := benchDoc(2000)
+	b.ReportAllocs()
+	b.SetBytes(int64(len(doc)))
+	for i := 0; i < b.N; i++ {
+		dec := yaml.NewDecoder(bytes.NewBufferString(doc))
+		dec.Origin(true, "bench.yaml")
+		var out any
+		if err := dec.Decode(&out); err != nil {
+			b.Fatal(err)
+		}
+	}
+}

yaml.go

@@ -144,6 +144,9 @@ func (dec *Decoder) Decode(v interface{}) (err error) {
 	d := newDecoder()
 	d.knownFields = dec.knownFields
 	d.origin = dec.origin
+	if d.origin {
+		d.origins = newOriginCache()
+	}
 	d.file = dec.file
 	d.disableTimestamps = dec.disableTimestamps
 	dec.parser.disableTimestamps = dec.disableTimestamps