diff --git a/ast/check.go b/ast/check.go index a2edbe66fa..217ab78856 100644 --- a/ast/check.go +++ b/ast/check.go @@ -269,7 +269,7 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) { } if tpe != nil { - env.tree.Put(path, tpe) + env.tree.Insert(path, tpe) } } diff --git a/ast/check_test.go b/ast/check_test.go index ea01767067..ef214d9392 100644 --- a/ast/check_test.go +++ b/ast/check_test.go @@ -351,6 +351,12 @@ func TestCheckInferenceRules(t *testing.T) { {`ref_rule_single_with_number_key`, `p.q[3] { true }`}, {`ref_regression_array_key`, `walker[[p, v]] = o { l = input; walk(l, k); [p, v] = k; o = {} }`}, + {`overlap`, `p.q[r] = y { x = ["a", "b"]; y = x[r] }`}, + {`overlap`, `p.q.r = false { true }`}, + {`overlap`, `p.q.r = "false" { true }`}, + {`overlap`, `p.q[42] = 1337 { true }`}, + {`overlap`, `p.q.a = input.a { true }`}, + {`overlap`, `p.q[56] = input.a { true }`}, } tests := []struct { @@ -504,6 +510,50 @@ func TestCheckInferenceRules(t *testing.T) { types.NewDynamicProperty(types.NewArray([]types.Type{types.NewArray(types.A, types.A), types.A}, nil), types.NewObject(nil, types.NewDynamicProperty(types.A, types.A))), )}, + { + note: "ref-rules single value, full ref to known leaf", + rules: ruleset2, + ref: "data.overlap.p.q.r", + expected: types.NewAny(types.B, types.S), + }, + { + note: "ref-rules single value, full ref to known leaf (same key type as dynamic, different value type)", + rules: ruleset2, + ref: "data.overlap.p.q[42]", + expected: types.N, + }, + { + note: "ref-rules single value, full ref to known leaf (any type)", + rules: ruleset2, + ref: "data.overlap.p.q.a", + expected: types.A, + }, + { + note: "ref-rules single value, full ref to known leaf (same key type as dynamic, any type)", + rules: ruleset2, + ref: "data.overlap.p.q[56]", + expected: types.A, + }, + { + note: "ref-rules single value, full ref to dynamic leaf", + rules: ruleset2, + ref: "data.overlap.p.q[1]", + expected: types.S, + }, + { + note: "ref-rules single value, prefix ref to partial object root", + rules: ruleset2, + ref: "data.overlap.p.q", + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty(json.Number("42"), types.N), + types.NewStaticProperty(json.Number("56"), types.A), + types.NewStaticProperty("a", types.A), + types.NewStaticProperty("r", types.Or(types.B, types.S)), + }, + types.NewDynamicProperty(types.N, types.S), + ), + }, } for _, tc := range tests { diff --git a/ast/env.go b/ast/env.go index e3a4e44284..d5a0706614 100644 --- a/ast/env.go +++ b/ast/env.go @@ -304,6 +304,116 @@ func (n *typeTreeNode) Put(path Ref, tpe types.Type) { curr.value = tpe } +// Insert inserts tpe at path in the tree, but also merges the value into any types.Object present along that path. +// If an types.Object is inserted, any leafs already present further down the tree are merged into the inserted object. +// path must be ground. +func (n *typeTreeNode) Insert(path Ref, tpe types.Type) { + curr := n + for i, term := range path { + c, ok := curr.children.Get(term.Value) + + var child *typeTreeNode + if !ok { + child = newTypeTree() + child.key = term.Value + curr.children.Put(child.key, child) + } else { + child = c.(*typeTreeNode) + + if child.value != nil && i+1 < len(path) { + // If child has an object value, merge the new value into it. + if o, ok := child.value.(*types.Object); ok { + var err error + child.value, err = insertIntoObject(o, path[i+1:], tpe) + if err != nil { + panic(fmt.Errorf("unreachable, insertIntoObject: %w", err)) + } + } + } + } + + curr = child + } + + curr.value = tpe + + if _, ok := tpe.(*types.Object); ok && curr.children.Len() > 0 { + // merge all leafs into the inserted object + leafs := curr.Leafs() + for p, t := range leafs { + var err error + curr.value, err = insertIntoObject(curr.value.(*types.Object), *p, t) + if err != nil { + panic(fmt.Errorf("unreachable, insertIntoObject: %w", err)) + } + } + } +} + +func insertIntoObject(o *types.Object, path Ref, tpe types.Type) (*types.Object, error) { + if len(path) == 0 { + return o, nil + } + + key, err := JSON(path[0].Value) + if err != nil { + return nil, fmt.Errorf("invalid path term %v: %w", path[0], err) + } + + if len(path) == 1 { + for _, prop := range o.StaticProperties() { + if util.Compare(prop.Key, key) == 0 { + prop.Value = types.Or(prop.Value, tpe) + return o, nil + } + } + staticProps := append(o.StaticProperties(), types.NewStaticProperty(key, tpe)) + return types.NewObject(staticProps, o.DynamicProperties()), nil + } + + for _, prop := range o.StaticProperties() { + if util.Compare(prop.Key, key) == 0 { + if propO := prop.Value.(*types.Object); propO != nil { + prop.Value, err = insertIntoObject(propO, path[1:], tpe) + if err != nil { + return nil, err + } + } else { + return nil, fmt.Errorf("cannot insert into non-object type %v", prop.Value) + } + return o, nil + } + } + + child, err := insertIntoObject(types.NewObject(nil, nil), path[1:], tpe) + if err != nil { + return nil, err + } + staticProps := append(o.StaticProperties(), types.NewStaticProperty(key, child)) + return types.NewObject(staticProps, o.DynamicProperties()), nil +} + +func (n *typeTreeNode) Leafs() map[*Ref]types.Type { + leafs := map[*Ref]types.Type{} + n.children.Iter(func(k, v util.T) bool { + collectLeafs(v.(*typeTreeNode), nil, leafs) + return false + }) + return leafs +} + +func collectLeafs(n *typeTreeNode, path Ref, leafs map[*Ref]types.Type) { + nPath := append(path, NewTerm(n.key)) + if n.Leaf() { + leafs[&nPath] = n.Value() + return + } + n.children.Iter(func(k, v util.T) bool { + collectLeafs(v.(*typeTreeNode), nPath, leafs) + return false + }) +} + func (n *typeTreeNode) Value() types.Type { return n.value } diff --git a/ast/env_test.go b/ast/env_test.go new file mode 100644 index 0000000000..1328b93de7 --- /dev/null +++ b/ast/env_test.go @@ -0,0 +1,285 @@ +// Copyright 2023 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package ast + +import ( + "encoding/json" + "testing" + + "github.com/open-policy-agent/opa/types" +) + +func TestInsertIntoObject(t *testing.T) { + tests := []struct { + note string + obj *types.Object + path Ref + tpe types.Type + expected types.Type + }{ + { + note: "adding to empty object", + obj: types.NewObject(nil, nil), + path: Ref{NewTerm(String("a"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + }, + { + note: "empty path", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: nil, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + }, + { + note: "adding to populated object", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(String("b"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.S), + }, + nil), + }, + { + note: "number key", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(Number("2"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty(json.Number("2"), types.S), + }, + nil), + }, + { + note: "same path, same type", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(String("a"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + }, + { + note: "same path, different type", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(String("a"))}, + tpe: types.B, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.Or(types.S, types.B))}, + nil), + }, + { + note: "same path, any type inserted", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(String("a"))}, + tpe: types.A, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.A)}, + nil), + }, + { + note: "same path, any type present", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.A)}, + nil), + path: Ref{NewTerm(String("a"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.A)}, + nil), + }, + { + note: "long path", + obj: types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("a", types.S)}, + nil), + path: Ref{NewTerm(String("b")), NewTerm(String("c")), NewTerm(String("d"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + }, + { + note: "long path, full match", + obj: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + path: Ref{NewTerm(String("b")), NewTerm(String("c")), NewTerm(String("d"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + }, + { + note: "long path, full match, different type", + obj: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + path: Ref{NewTerm(String("b")), NewTerm(String("c")), NewTerm(String("d"))}, + tpe: types.B, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.Or(types.S, types.B)), + }, nil)), + }, nil)), + }, + nil), + }, + { + note: "long path, partial match", + obj: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + path: Ref{NewTerm(String("b")), NewTerm(String("x")), NewTerm(String("d"))}, + tpe: types.S, + expected: types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("a", types.S), + types.NewStaticProperty("b", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("c", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + types.NewStaticProperty("x", types.NewObject([]*types.StaticProperty{ + types.NewStaticProperty("d", types.S), + }, nil)), + }, nil)), + }, + nil), + }, + } + + for _, tc := range tests { + t.Run(tc.note, func(t *testing.T) { + result, err := insertIntoObject(tc.obj, tc.path, tc.tpe) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if types.Compare(result, tc.expected) != 0 { + t.Fatalf("Expected %v but got %v", tc.expected, result) + } + }) + } +} + +func TestTypeTreeInsert(t *testing.T) { + n := newTypeTree() + + abcRef := Ref{NewTerm(String("a")), NewTerm(String("b")), NewTerm(String("c"))} + n.Put(abcRef, types.B) + actual := n.Get(abcRef) + if types.Compare(actual, types.B) != 0 { + t.Fatalf("Expected %v but got %v", types.B, actual) + } + + abdeRef := Ref{NewTerm(String("a")), NewTerm(String("b")), NewTerm(String("d")), NewTerm(String("e"))} + n.Put(abdeRef, types.N) + actual = n.Get(abdeRef) + if types.Compare(actual, types.N) != 0 { + t.Fatalf("Expected %v but got %v", types.N, actual) + } + + // existing "child" leafs should be added to new intermediate object leaf + + abRef := Ref{NewTerm(String("a")), NewTerm(String("b"))} + n.Insert(abRef, types.NewObject(nil, &types.DynamicProperty{Key: types.N, Value: types.S})) + + actual = n.Get(abRef) + expected := types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("c", types.B), + types.NewStaticProperty("d", types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("e", types.N)}, + nil)), + }, + &types.DynamicProperty{Key: types.N, Value: types.S}, + ) + if types.Compare(actual, expected) != 0 { + t.Fatalf("Expected %v but got %v", expected, actual) + } + + // new "child" leafs should be added to new intermediate object leaf + + abfRef := Ref{NewTerm(String("a")), NewTerm(String("b")), NewTerm(String("f"))} + n.Insert(abfRef, types.S) + + actual = n.Get(abfRef) + if types.Compare(actual, types.S) != 0 { + t.Fatalf("Expected %v but got %v", types.S, actual) + } + + actual = n.Get(abRef) + expected = types.NewObject( + []*types.StaticProperty{ + types.NewStaticProperty("c", types.B), + types.NewStaticProperty("f", types.S), + types.NewStaticProperty("d", types.NewObject( + []*types.StaticProperty{types.NewStaticProperty("e", types.N)}, + nil)), + }, + &types.DynamicProperty{Key: types.N, Value: types.S}, + ) + if types.Compare(actual, expected) != 0 { + t.Fatalf("Expected %v but got %v", expected, actual) + } +} diff --git a/test/cases/testdata/type/test-regressions.yaml b/test/cases/testdata/type/test-regressions.yaml new file mode 100644 index 0000000000..c0c51b83bf --- /dev/null +++ b/test/cases/testdata/type/test-regressions.yaml @@ -0,0 +1,32 @@ +cases: +- note: regression/partial-object override, different key type, query + modules: + - | + package test + + p[k] := v { + v := ["a", "b", "c"][k] + } + + p.foo := "bar" + query: data.test.p.foo = x + want_result: + - x: "bar" +- note: regression/partial-object override, different key type, referenced in other rule + modules: + - | + package test + + p[k] := v { + v := ["a", "b", "c"][k] + } + + p.foo := "bar" + + q[x] { + x := p[_] + x == "bar" + } + query: data.test.q = x + want_result: + - x: ["bar"] \ No newline at end of file diff --git a/types/types_test.go b/types/types_test.go index 9d0f4827e2..3f58ef77e9 100644 --- a/types/types_test.go +++ b/types/types_test.go @@ -242,6 +242,8 @@ func TestSelect(t *testing.T) { {"non int-2", NewArray([]Type{S, N, B}, nil), 1, nil}, {"static", NewObject([]*StaticProperty{NewStaticProperty("hello", S)}, nil), "hello", S}, {"dynamic", NewObject([]*StaticProperty{NewStaticProperty("hello", S)}, NewDynamicProperty(S, N)), "goodbye", N}, + {"dynamic, different key types", NewObject([]*StaticProperty{NewStaticProperty("hello", S)}, NewDynamicProperty(N, N)), json.Number("2"), N}, + {"dynamic, different key types", NewObject([]*StaticProperty{NewStaticProperty("hello", S)}, NewDynamicProperty(N, N)), "hello", S}, {"non exist", NewObject([]*StaticProperty{NewStaticProperty("hello", S)}, nil), "deadbeef", nil}, {"non string", NewObject([]*StaticProperty{NewStaticProperty(json.Number("1"), S), NewStaticProperty(json.Number("2"), N)}, nil), json.Number("2"), N}, {"member of", NewSet(N), json.Number("2"), N},