diff --git a/cli/test.yaml b/cli/test.yaml index a2ef24f..c669f75 100644 --- a/cli/test.yaml +++ b/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)' diff --git a/compiler.go b/compiler.go index 050951b..7747077 100644 --- a/compiler.go +++ b/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( - &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}) @@ -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( - &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 ( + // 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"}}, - &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) @@ -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"}}, - &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}, @@ -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: 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..6f7c1f8 100644 --- a/compiler_test.go +++ b/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 { diff --git a/inplace.go b/inplace.go new file mode 100644 index 0000000..442babb --- /dev/null +++ b/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() +}