perf: update the accumulated value of a reduce expression in place

_assign and _modify create an allocator on every call, so the value they
are given is never one that allocator created, and updateObject and
updateArrayIndex always take the copying branch. Accumulating into an
object or an array with a reduce expression therefore copies the whole
accumulated value on every iteration, which is quadratic in the number
of elements. The query of #241,

    [range(20000)] | reduce .[] as $item ({}; .[($item|tostring)] = {})

took between 10 and 21 seconds here over three runs, against 0.03
seconds for jq.

Create the allocator once per reduce expression and pass it to the
assignment as a third argument, so that the accumulated value is
allocated on the first iteration and modified in place afterwards. The
query above now takes 0.03 seconds, and the memory allocated by an
accumulating reduction grows linearly instead of quadratically: 2.1 MB
instead of 545 MB for 4000 elements.

This is only correct while the accumulated value, and the values the
allocator created inside it, are reachable through the reduction alone.
Query#inPlaceUpdate therefore requires the update query to be a single
assignment, and its right-hand side to neither observe the input value,
which would let the assigned value alias the value it is assigned into,
nor yield more than one value, which would run the assignment again
against a value it already modified. Everything else keeps copying. A
foreach expression is left alone, because it yields every intermediate
accumulated value.

cli/test.yaml

@@ -3642,6 +3642,68 @@
   expected: |
     [[1,1],[1]]
 
+- name: reduce with assignment in update
+  args:
+    - -c
+    - 'reduce range(3) as $i ({}; .[$i|tostring] = {}),
+       reduce range(3) as $i ({}; .a[$i|tostring] = $i),
+       reduce range(3) as $i ([]; .[$i] = $i),
+       reduce range(3) as $i ({}; .[$i|tostring] += 1)'
+  input: 'null'
+  expected: |
+    {"0":{},"1":{},"2":{}}
+    {"a":{"0":0,"1":1,"2":2}}
+    [0,1,2]
+    {"0":1,"1":1,"2":1}
+
+- name: reduce with assignment of the accumulated value
+  args:
+    - -c
+    - 'reduce range(3) as $i ({}; .a = .),
+       reduce range(3) as $i ({}; .a = {v: .}),
+       reduce range(3) as $i ({}; . as $p | .a = $i | .b = $p)'
+  input: 'null'
+  expected: |
+    {"a":{"a":{"a":{}}}}
+    {"a":{"v":{"a":{"v":{"a":{"v":{}}}}}}}
+    {"a":2,"b":{"a":1,"b":{"a":0,"b":{}}}}
+
+- name: reduce with assignment of multiple values
+  args:
+    - -c
+    - 'reduce range(2) as $i ({"k":"k"}; .[.k] = ("x","y")),
+       reduce range(2) as $i ({}; (.a = 1), (.b = 2))'
+  input: 'null'
+  expected: |
+    {"k":"y","y":"y"}
+    {"b":2}
+
+- name: reduce with assignment of a shared value
+  args:
+    - -c
+    - '{"a":1} as $x | reduce range(2) as $i ({}; .[$i|tostring] = $x) |
+       [(.["0"].a = 9), $x]'
+  input: 'null'
+  expected: |
+    [{"0":{"a":9},"1":{"a":1}},{"a":1}]
+
+- name: reduce with assignment result reused
+  args:
+    - -c
+    - 'reduce range(3) as $i ({}; .[$i|tostring] = $i) as $r |
+       [($r | .["0"] = 9), $r]'
+  input: 'null'
+  expected: |
+    [{"0":9,"1":1,"2":2},{"0":0,"1":1,"2":2}]
+
+- name: foreach with assignment in update
+  args:
+    - -c
+    - '[foreach range(3) as $i ({}; .[$i|tostring] = $i)]'
+  input: 'null'
+  expected: |
+    [{"0":0},{"0":0,"1":1},{"0":0,"1":1,"2":2}]
+
 - name: foreach iterator
   args:
     - 'foreach range(5) as $item (0; $item)'

compiler.go

@@ -393,7 +393,7 @@ func (c *compiler) compileQuery(e *Query) error {
 		return c.compileAlt(e.Left, e.Right)
 	case OpAssign, OpModify, OpUpdateAdd, OpUpdateSub,
 		OpUpdateMul, OpUpdateDiv, OpUpdateMod, OpUpdateAlt:
-		return c.compileQueryUpdate(e.Left, e.Right, e.Op)
+		return c.compileQueryUpdate(e.Left, e.Right, e.Op, nil)
 	case OpOr:
 		return c.compileIf(
 			&If{
@@ -467,7 +467,10 @@ func (c *compiler) compileAlt(l, r *Query) error {
 	return c.compileQuery(r)
 }
 
-func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
+// compileQueryUpdate compiles an assignment operator. alloc is nil unless the
+// assignment may update the input value in place, in which case it yields the
+// allocator of the enclosing reduce expression.
+func (c *compiler) compileQueryUpdate(l, r *Query, op Operator, alloc *Query) error {
 	switch op {
 	case OpAssign:
 		// optimize assignment operator with constant indexing and slicing
@@ -487,12 +490,7 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
 		}
 		fallthrough
 	case OpModify:
-		return c.compileFunc(
+		return c.compileUpdateFunc(op.getFunc(), []*Query{l, r}, alloc)
-			&Func{
-				Name: op.getFunc(),
-				Args: []*Query{l, r},
-			},
-		)
 	default:
 		name := "$%0"
 		c.append(&code{op: opdup})
@@ -500,27 +498,53 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
 			return err
 		}
 		c.append(&code{op: opstore, v: c.pushVariable(name)})
-		return c.compileFunc(
+		return c.compileUpdateFunc(
-			&Func{
+			"_modify",
-				Name: "_modify",
+			[]*Query{
-				Args: []*Query{
+				l,
-					l,
+				{Term: &Term{
-					{Term: &Term{
+					Type: TermTypeFunc,
-						Type: TermTypeFunc,
+					Func: &Func{
-						Func: &Func{
+						Name: op.getFunc(),
-							Name: op.getFunc(),
+						Args: []*Query{
-							Args: []*Query{
+							{Term: &Term{Type: TermTypeIdentity}},
-								{Term: &Term{Type: TermTypeIdentity}},
+							{Term: &Term{Type: TermTypeFunc, Func: &Func{Name: name}}},
-								{Term: &Term{Type: TermTypeFunc, Func: &Func{Name: name}}},
-							},
 						},
-					}},
+					},
-				},
+				}},
 			},
+			alloc,
 		)
 	}
 }
 
+const (
+	// Both names are deliberately invalid identifiers, so that a query cannot
+	// reach the allocator variants directly.
+	allocatorArgSuffix = "/3"
+	allocatorName      = "$%allocator"
+)
+
+// compileUpdateFunc calls _assign or _modify, or the variant taking an
+// allocator as its third argument.
+func (c *compiler) compileUpdateFunc(name string, args []*Query, alloc *Query) error {
+	if alloc == nil {
+		return c.compileFunc(&Func{Name: name, Args: args})
+	}
+	args = append(args, alloc)
+	fn := c.lookupBuiltin(name+allocatorArgSuffix, len(args))
+	if fn == nil {
+		switch name {
+		case "_assign":
+			c.compileAssign(true)
+		case "_modify":
+			c.compileModify(true)
+		}
+		fn = c.lookupBuiltin(name+allocatorArgSuffix, len(args))
+	}
+	return c.compileCallPc(fn, args)
+}
+
 func (c *compiler) compileBind(l, r *Query, patterns []*Pattern) error {
 	defer c.newScopeDepth()()
 	c.append(&code{op: opdup})
@@ -716,6 +740,21 @@ func (c *compiler) compileReduce(e *Reduce) error {
 	}
 	f()
 	c.append(&code{op: opstore, v: v})
+	// An assignment forming the whole update query is the only thing that
+	// observes the accumulated value, so it may modify it in place; see
+	// Query#inPlaceUpdate. The allocator lives across the whole reduction.
+	l, r, op, inPlace := e.Update.inPlaceUpdate()
+	var alloc *Query
+	if inPlace {
+		c.appends(
+			&code{op: oppush, v: nil},
+			&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
+			&code{op: opstore, v: c.pushVariable(allocatorName)},
+		)
+		alloc = &Query{Term: &Term{
+			Type: TermTypeFunc, Func: &Func{Name: allocatorName},
+		}}
+	}
 	setfork := c.lazy(func() *code {
 		return &code{op: opfork, v: len(c.codes)}
 	})
@@ -727,7 +766,13 @@ func (c *compiler) compileReduce(e *Reduce) error {
 	}
 	c.append(&code{op: opload, v: v})
 	f = c.newScopeDepth()
-	if err := c.compileQuery(e.Update); err != nil {
+	var err error
+	if inPlace {
+		err = c.compileQueryUpdate(l, r, op, alloc)
+	} else {
+		err = c.compileQuery(e.Update)
+	}
+	if err != nil {
 		return err
 	}
 	f()
@@ -932,9 +977,9 @@ func (c *compiler) compileFunc(e *Func) error {
 		if !compiled {
 			switch e.Name {
 			case "_assign":
-				c.compileAssign()
+				c.compileAssign(false)
 			case "_modify":
-				c.compileModify()
+				c.compileModify(false)
 			case "_last":
 				c.compileLast()
 			}
@@ -1030,29 +1075,54 @@ func (c *compiler) compileFunc(e *Func) error {
 //
 // To overcome the difficulty of reducing allocations on `setpath`, we use the
 // `allocator` type and track the allocated addresses during the reduction.
-func (c *compiler) compileAssign() {
+//
-	defer c.appendBuiltin("_assign", 2)()
+// The shared variant reuses the allocator of the enclosing reduce expression
+// instead of creating a new one, modifying the value in place.
+func (c *compiler) compileAssign(shared bool) {
+	name, argcnt := "_assign", 2
+	if shared {
+		name, argcnt = name+allocatorArgSuffix, 3
+	}
+	defer c.appendBuiltin(name, argcnt)()
 	scope := c.newScope()
 	v, p := [2]int{scope.id, 0}, [2]int{scope.id, 1}
 	x, a := [2]int{scope.id, 2}, [2]int{scope.id, 3}
 	// Cannot reuse v, p due to backtracking in x.
 	w, q := [2]int{scope.id, 4}, [2]int{scope.id, 5}
 	c.appends(
-		&code{op: opscope, v: [3]int{scope.id, 6, 2}},
+		&code{op: opscope, v: [3]int{scope.id, 6, argcnt}},
 		&code{op: opstore, v: v}, //                def _assign(p; $x):
 		&code{op: opstore, v: p},
 		&code{op: opstore, v: x},
+	)
+	if shared {
+		c.append(&code{op: opstore, v: a}) //       the allocator argument
+	}
+	c.appends(
 		&code{op: opload, v: v},
 		&code{op: opexpbegin},
 		&code{op: opload, v: x},
 		&code{op: opcallpc},
 		&code{op: opstore, v: x},
 		&code{op: opexpend},
-		&code{op: oppush, v: nil},
+	)
-		&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
+	if shared {
-		&code{op: opstore, v: a},
+		c.appends(
+			&code{op: opload, v: v},
+			&code{op: opload, v: a},
+			&code{op: opcallpc},
+			&code{op: opstore, v: a},
+		)
+	} else {
+		c.appends(
+			&code{op: oppush, v: nil},
+			&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
+			&code{op: opstore, v: a},
+		)
+	}
+	c.appends(
 		&code{op: opload, v: v},
-		&code{op: opfork, v: len(c.codes) + 30}, // reduce [L1]
+		&code{op: opfork, v: len(c.codes) + 17}, // reduce [L1]
 		&code{op: opdup},
 		&code{op: opstore, v: w},
 		&code{op: oppathbegin}, //                  path(p)
@@ -1076,24 +1146,49 @@ func (c *compiler) compileAssign() {
 
 // Appends the compiled code for the update-assignment operator (`|=`) to
 // maximize performance. We use the `allocator` type, just like `_assign/2`.
-func (c *compiler) compileModify() {
+//
-	defer c.appendBuiltin("_modify", 2)()
+// The shared variant reuses the allocator of the enclosing reduce expression;
+// see compiler#compileAssign.
+func (c *compiler) compileModify(shared bool) {
+	name, argcnt := "_modify", 2
+	if shared {
+		name, argcnt = name+allocatorArgSuffix, 3
+	}
+	defer c.appendBuiltin(name, argcnt)()
 	scope := c.newScope()
 	v, p := [2]int{scope.id, 0}, [2]int{scope.id, 1}
 	f, d := [2]int{scope.id, 2}, [2]int{scope.id, 3}
 	a, l := [2]int{scope.id, 4}, [2]int{scope.id, 5}
 	c.appends(
-		&code{op: opscope, v: [3]int{scope.id, 6, 2}},
+		&code{op: opscope, v: [3]int{scope.id, 6, argcnt}},
 		&code{op: opstore, v: v}, //                def _modify(p; f):
 		&code{op: opstore, v: p},
 		&code{op: opstore, v: f},
+	)
+	if shared {
+		c.append(&code{op: opstore, v: a}) //       the allocator argument
+	}
+	c.appends(
 		&code{op: oppush, v: []any{}},
 		&code{op: opstore, v: d},
-		&code{op: oppush, v: nil},
+	)
-		&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
+	if shared {
-		&code{op: opstore, v: a},
+		c.appends(
+			&code{op: opload, v: v},
+			&code{op: opload, v: a},
+			&code{op: opcallpc},
+			&code{op: opstore, v: a},
+		)
+	} else {
+		c.appends(
+			&code{op: oppush, v: nil},
+			&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
+			&code{op: opstore, v: a},
+		)
+	}
+	c.appends(
 		&code{op: opload, v: v},
-		&code{op: opfork, v: len(c.codes) + 40}, // reduce [L1]
+		&code{op: opfork, v: len(c.codes) + 31}, // reduce [L1]
 		&code{op: oppathbegin},                  // path(p)
 		&code{op: opload, v: p},
 		&code{op: opcallpc},
@@ -1102,7 +1197,7 @@ func (c *compiler) compileModify() {
 		&code{op: opstore, v: p},                // as $p (.;
 		&code{op: opforklabel, v: l},            // label $l |
 		&code{op: opload, v: v},                 //
-		&code{op: opfork, v: len(c.codes) + 37}, // [L2]
+		&code{op: opfork, v: len(c.codes) + 28}, // [L2]
 		&code{op: oppop},                        // (getpath($p) |
 		&code{op: opload, v: a},
 		&code{op: opload, v: a},
@@ -1116,8 +1211,8 @@ func (c *compiler) compileModify() {
 		&code{op: opcall, v: [3]any{funcSetpathWithAllocator, 3, "_setpath"}},
 		&code{op: opstore, v: v},
 		&code{op: opload, v: v},                 // ., break $l
-		&code{op: opfork, v: len(c.codes) + 35}, // [L4]
+		&code{op: opfork, v: len(c.codes) + 26}, // [L4]
-		&code{op: opjump, v: len(c.codes) + 39}, // [L3]
+		&code{op: opjump, v: len(c.codes) + 30}, // [L3]
 		&code{op: opload, v: l},                 // [L4]
 		&code{op: opcall, v: [3]any{funcBreak(""), 0, "_break"}},
 		&code{op: opload, v: p},   //               append $p to $d [L2]

compiler_test.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	"os"
 	"reflect"
+	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -374,6 +375,57 @@ func TestCodeRun_RaceRegexp(t *testing.T) {
 	wg.Wait()
 }
 
+// Building an object or an array of n elements must allocate memory
+// proportional to n, not to n squared.
+func TestCodeRun_ReduceAssignmentAllocation(t *testing.T) {
+	for _, src := range []string{
+		`reduce .[] as $x ({}; .[$x|tostring] = $x)`,
+		`reduce .[] as $x ({}; .[$x|tostring] |= $x)`,
+		`reduce .[] as $x ({}; .[$x|tostring] += 1)`,
+		`reduce .[] as $x ([]; .[$x] = $x)`,
+	} {
+		t.Run(src, func(t *testing.T) {
+			query, err := gojq.Parse(src)
+			if err != nil {
+				t.Fatal(err)
+			}
+			code, err := gojq.Compile(query)
+			if err != nil {
+				t.Fatal(err)
+			}
+			small, large := allocatedBytes(t, code, 1000), allocatedBytes(t, code, 4000)
+			// Quadratic growth would be about sixteen times as much.
+			if expected := small * 8; large > expected {
+				t.Errorf("expected at most %v bytes for 4000 elements, got %v (%v bytes for 1000 elements)",
+					expected, large, small)
+			}
+		})
+	}
+}
+
+func allocatedBytes(t *testing.T, code *gojq.Code, n int) uint64 {
+	t.Helper()
+	input := make([]any, n)
+	for i := range input {
+		input[i] = i
+	}
+	var before, after runtime.MemStats
+	runtime.GC()
+	runtime.ReadMemStats(&before)
+	iter := code.Run(input)
+	for {
+		v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if err, ok := v.(error); ok {
+			t.Fatal(err)
+		}
+	}
+	runtime.ReadMemStats(&after)
+	return after.TotalAlloc - before.TotalAlloc
+}
+
 func BenchmarkCompile(b *testing.B) {
 	cnt, err := os.ReadFile("builtin.jq")
 	if err != nil {

inplace.go

@@ -0,0 +1,204 @@
+package gojq
+
+// Static analysis deciding whether the accumulated value of a reduce
+// expression may be updated in place; see compiler#compileReduce. The checks
+// are conservative: failing them only keeps the previous copying behavior.
+
+// inPlaceUpdate reports whether the query is a single assignment applied to
+// the input value, and returns its operands.
+func (e *Query) inPlaceUpdate() (l, r *Query, op Operator, ok bool) {
+	if e.Term != nil || len(e.FuncDefs) > 0 {
+		return nil, nil, 0, false
+	}
+	switch e.Op {
+	case OpModify:
+		// _modify releases the allocator before calling the update function.
+		return e.Left, e.Right, e.Op, true
+	case OpAssign:
+		// A constant path compiles to setpath, which uses no allocator.
+		if e.Left.toIndices(nil) != nil {
+			return nil, nil, 0, false
+		}
+		fallthrough
+	case OpUpdateAdd, OpUpdateSub, OpUpdateMul,
+		OpUpdateDiv, OpUpdateMod, OpUpdateAlt:
+		// The right-hand side must neither observe the input value, which it
+		// could otherwise alias, nor yield multiple values, which would run
+		// the assignment against an already modified input.
+		if !e.Right.readsInput() && e.Right.singleValued() {
+			return e.Left, e.Right, e.Op, true
+		}
+	}
+	return nil, nil, 0, false
+}
+
+// readsInput reports whether evaluating the query may observe its input value.
+// It answers true unless it can prove otherwise.
+func (e *Query) readsInput() bool {
+	if len(e.FuncDefs) > 0 {
+		return true
+	}
+	if e.Term != nil {
+		return e.Term.readsInput()
+	}
+	switch e.Op {
+	case OpPipe:
+		// Binding variables also exposes the input to the right-hand side.
+		if len(e.Patterns) > 0 {
+			return e.Left.readsInput() || e.Right.readsInput()
+		}
+		return e.Left.readsInput()
+	case OpComma, OpAdd, OpSub, OpMul, OpDiv, OpMod, OpEq, OpNe,
+		OpGt, OpLt, OpGe, OpLe, OpAnd, OpOr, OpAlt:
+		return e.Left.readsInput() || e.Right.readsInput()
+	default:
+		return true
+	}
+}
+
+// singleValued reports whether the query yields at most one value.
+// It answers false unless it can prove otherwise.
+func (e *Query) singleValued() bool {
+	if len(e.FuncDefs) > 0 {
+		return false
+	}
+	if e.Term != nil {
+		return e.Term.singleValued()
+	}
+	switch e.Op {
+	case OpPipe:
+		return len(e.Patterns) == 0 &&
+			e.Left.singleValued() && e.Right.singleValued()
+	case OpAdd, OpSub, OpMul, OpDiv, OpMod, OpEq, OpNe,
+		OpGt, OpLt, OpGe, OpLe, OpAnd, OpOr:
+		return e.Left.singleValued() && e.Right.singleValued()
+	default:
+		return false
+	}
+}
+
+func (e *Term) readsInput() bool {
+	for _, s := range e.SuffixList {
+		if s.readsInput() {
+			return true
+		}
+	}
+	switch e.Type {
+	case TermTypeNull, TermTypeTrue, TermTypeFalse, TermTypeNumber:
+		return false
+	case TermTypeString:
+		return e.Str.readsInput()
+	case TermTypeObject:
+		for _, kv := range e.Object.KeyVals {
+			if kv.readsInput() {
+				return true
+			}
+		}
+		return false
+	case TermTypeArray:
+		return e.Array.Query != nil && e.Array.Query.readsInput()
+	case TermTypeUnary:
+		return e.Unary.Term.readsInput()
+	case TermTypeFunc:
+		// Any function may be redefined, so only $x is known not to read it.
+		return e.Func.Name[0] != '$' || len(e.Func.Args) > 0
+	case TermTypeQuery:
+		return e.Query.readsInput()
+	default:
+		return true
+	}
+}
+
+func (e *Term) singleValued() bool {
+	for _, s := range e.SuffixList {
+		if !s.singleValued() {
+			return false
+		}
+	}
+	switch e.Type {
+	case TermTypeNull, TermTypeTrue, TermTypeFalse, TermTypeNumber:
+		return true
+	case TermTypeString:
+		return e.Str.singleValued()
+	case TermTypeObject:
+		for _, kv := range e.Object.KeyVals {
+			if !kv.singleValued() {
+				return false
+			}
+		}
+		return true
+	case TermTypeArray:
+		return true
+	case TermTypeUnary:
+		return e.Unary.Term.singleValued()
+	case TermTypeFunc:
+		return e.Func.Name[0] == '$' && len(e.Func.Args) == 0
+	case TermTypeQuery:
+		return e.Query.singleValued()
+	default:
+		return false
+	}
+}
+
+func (e *Suffix) readsInput() bool {
+	// Only an index is evaluated against the input value.
+	return e.Index != nil && e.Index.readsInput()
+}
+
+func (e *Suffix) singleValued() bool {
+	return !e.Iter && (e.Index == nil || e.Index.singleValued())
+}
+
+func (e *Index) readsInput() bool {
+	if e.Str != nil {
+		return e.Str.readsInput()
+	}
+	return e.Start != nil && e.Start.readsInput() ||
+		e.End != nil && e.End.readsInput()
+}
+
+func (e *Index) singleValued() bool {
+	if e.Str != nil {
+		return e.Str.singleValued()
+	}
+	return (e.Start == nil || e.Start.singleValued()) &&
+		(e.End == nil || e.End.singleValued())
+}
+
+func (e *String) readsInput() bool {
+	for _, q := range e.Queries {
+		if q.readsInput() {
+			return true
+		}
+	}
+	return false
+}
+
+func (e *String) singleValued() bool {
+	for _, q := range e.Queries {
+		if !q.singleValued() {
+			return false
+		}
+	}
+	return true
+}
+
+func (e *ObjectKeyVal) readsInput() bool {
+	if e.Val == nil {
+		// {foo} reads the input value, {$foo} does not.
+		return e.Key == "" || e.Key[0] != '$'
+	}
+	if e.KeyString != nil && e.KeyString.readsInput() ||
+		e.KeyQuery != nil && e.KeyQuery.readsInput() {
+		return true
+	}
+	return e.Val.readsInput()
+}
+
+func (e *ObjectKeyVal) singleValued() bool {
+	if e.KeyString != nil && !e.KeyString.singleValued() ||
+		e.KeyQuery != nil && !e.KeyQuery.singleValued() {
+		return false
+	}
+	return e.Val == nil || e.Val.singleValued()
+}