diff --git a/cli/test.yaml b/cli/test.yaml index 809d68c..d27616f 100644 --- a/cli/test.yaml +++ b/cli/test.yaml @@ -3609,6 +3609,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)' diff --git a/compiler.go b/compiler.go index 477650c..978d2fa 100644 --- a/compiler.go +++ b/compiler.go @@ -397,7 +397,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{ @@ -471,7 +471,11 @@ func (c *compiler) compileAlt(l, r *Query) error { return c.compileQuery(r) } -func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error { +// compileQueryUpdate appends the compiled code for an assignment operator. The +// alloc query is nil unless the assignment may update the input value in +// place, in which case it yields the allocator of the enclosing reduce +// expression; see compiler#compileReduce. +func (c *compiler) compileQueryUpdate(l, r *Query, op Operator, alloc *Query) error { switch op { case OpAssign: // optimize assignment operator with constant indexing and slicing @@ -491,12 +495,7 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error { } fallthrough case OpModify: - return c.compileFunc( - &Func{ - Name: op.getFunc(), - Args: []*Query{l, r}, - }, - ) + return c.compileUpdateFunc(op.getFunc(), []*Query{l, r}, alloc) default: name := "$%0" c.append(&code{op: opdup}) @@ -504,27 +503,59 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error { return err } c.append(&code{op: opstore, v: c.pushVariable(name)}) - return c.compileFunc( - &Func{ - Name: "_modify", - Args: []*Query{ - l, - {Term: &Term{ - Type: TermTypeFunc, - Func: &Func{ - Name: op.getFunc(), - Args: []*Query{ - {Term: &Term{Type: TermTypeIdentity}}, - {Term: &Term{Type: TermTypeFunc, Func: &Func{Name: name}}}, - }, + return c.compileUpdateFunc( + "_modify", + []*Query{ + l, + {Term: &Term{ + Type: TermTypeFunc, + Func: &Func{ + Name: op.getFunc(), + Args: []*Query{ + {Term: &Term{Type: TermTypeIdentity}}, + {Term: &Term{Type: TermTypeFunc, Func: &Func{Name: name}}}, }, - }}, - }, + }, + }}, }, + alloc, ) } } +const ( + // Appended to the names of the _assign and _modify variants taking an + // allocator as an additional argument. The resulting name is not a valid + // function name, so that a query cannot call such a variant directly with + // an arbitrary value in place of the allocator. + allocatorArgSuffix = "/3" + + // The name of the variable holding the allocator of a reduce expression. + // It is not a valid variable name, just like the variable of the operand + // of an update-assignment operator. + allocatorName = "$%allocator" +) + +// compileUpdateFunc appends a call to _assign or _modify, or to the variant +// taking the allocator to update the input value with 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}) @@ -720,6 +751,25 @@ func (c *compiler) compileReduce(e *Reduce) error { } f() c.append(&code{op: opstore, v: v}) + // The accumulated value is passed to the update query and replaced by its + // result, so an assignment forming the whole update query is the only + // thing that observes it. Such an assignment may modify the accumulated + // value in place, instead of copying it on every iteration, as long as it + // keeps track of the values it allocated; see Query#inPlaceUpdate. The + // allocator is created once, outside the loop, and lives until the + // reduction finishes. + 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)} }) @@ -731,7 +781,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() @@ -936,9 +992,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() } @@ -1034,29 +1090,56 @@ 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 takes the allocator of the enclosing reduce expression as +// its third argument instead of creating a new one, so that the value being +// assigned to is modified in place on every iteration but the first; see +// compiler#compileReduce. +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"}}, - &code{op: opstore, v: a}, + ) + if shared { + 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) @@ -1080,24 +1163,50 @@ 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 takes the allocator of the enclosing reduce expression as +// its third argument instead of creating a new one; see +// compiler#compileReduce and 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"}}, - &code{op: opstore, v: a}, + ) + if shared { + 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}, @@ -1106,7 +1215,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}, @@ -1120,8 +1229,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: opjump, v: len(c.codes) + 39}, // [L3] + &code{op: opfork, v: len(c.codes) + 26}, // [L4] + &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] diff --git a/compiler_test.go b/compiler_test.go index fe67bf5..797a8c0 100644 --- a/compiler_test.go +++ b/compiler_test.go @@ -7,6 +7,7 @@ import ( "log" "os" "reflect" + "runtime" "sync" "testing" "time" @@ -374,6 +375,59 @@ func TestCodeRun_RaceRegexp(t *testing.T) { wg.Wait() } +// An assignment forming the whole update query of a reduce expression modifies +// the accumulated value in place, so building an object or an array of n +// elements allocates 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) + // Copying the accumulated value on every iteration allocates about + // sixteen times as much memory for four times as many elements. + 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 { diff --git a/inplace.go b/inplace.go new file mode 100644 index 0000000..6a401e7 --- /dev/null +++ b/inplace.go @@ -0,0 +1,225 @@ +package gojq + +// This file implements the static analysis deciding whether the accumulated +// value of a reduce expression may be updated in place; see +// compiler#compileReduce for how the result is used, and the comment on +// compiler#compileAssign for why an allocator makes assignments cheap. +// +// Reusing an allocator across the iterations of a reduce expression is only +// correct while the accumulated value, and the values the allocator created +// inside it, are reachable through the accumulated value alone. The checks +// below are deliberately conservative; failing them only means the assignment +// keeps copying the accumulated value, which is the previous behavior. + +// inPlaceUpdate reports whether the query is a single assignment applied to +// the input value, and returns its operands. Such an assignment receives the +// accumulated value of the enclosing reduce expression, and passes it on as +// the next accumulated value, without any other part of the query observing +// it in between. +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: + // The update function is applied to the values found at the target + // paths, so it never observes the input value itself. When a target + // path points to a value the allocator created, _modify releases the + // allocator before calling the update function. + return e.Left, e.Right, e.Op, true + case OpAssign: + // An assignment to a constant path is compiled to a setpath call, + // which does not use an allocator; see compiler#compileQueryUpdate. + if e.Left.toIndices(nil) != nil { + return nil, nil, 0, false + } + fallthrough + case OpUpdateAdd, OpUpdateSub, OpUpdateMul, + OpUpdateDiv, OpUpdateMod, OpUpdateAlt: + // The right-hand side is evaluated against the input value, so it must + // not observe it; the assigned value would otherwise be able to alias + // the value it is assigned into. It must not yield multiple values + // either, because the assignment then runs once per value, each time + // against an input value the previous run already modified. + 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: + // The right-hand side is evaluated against the output of the + // left-hand side, unless the query binds variables, in which case it + // is evaluated against the input value as well. + 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: + // Only a variable reference is known not to read the input value; any + // function may be redefined, so its name tells us nothing. + 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 { + // Both iterating and the optional operator work on the value of the term, + // while 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 { + // The shorthand {foo} and {"foo"} read the input value, while {$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() +}