ast: Making partial object key rules contribute to dynamic portion of object type (#6177)

... instead of static portion.

This change allows object unification/comparison between an object composed at eval-time (i.e. constructed by rules) and a static object that doesn't contain all keys declared by rules, as those might not be defined at eval-time.

Fixes: #6138
Signed-off-by: Johan Fylling <johan.dev@fylling.se>

ast/check.go

@@ -268,7 +268,7 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) {
 	}
 
 	if tpe != nil {
-		env.tree.Insert(path, tpe)
+		env.tree.Insert(path, tpe, env)
 	}
 }
 

ast/check_test.go

@@ -355,8 +355,8 @@ func TestCheckInferenceRules(t *testing.T) {
 		{`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 }`},
+		{`overlap`, `p.q2.a = input.a { true }`},
+		{`overlap`, `p.q2[56] = input.a { true }`},
 	}
 
 	tests := []struct {
@@ -525,33 +525,28 @@ func TestCheckInferenceRules(t *testing.T) {
 		{
 			note:     "ref-rules single value, full ref to known leaf (any type)",
 			rules:    ruleset2,
-			ref:      "data.overlap.p.q.a",
+			ref:      "data.overlap.p.q2.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]",
+			ref:      "data.overlap.p.q2[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,
+			expected: types.Any{types.B, types.N, types.S}, // key type cannot be tied to specific dynamic value type, so we get all of them
 		},
 		{
 			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),
+				nil,
+				types.NewDynamicProperty(types.Any{types.N, types.S}, types.Any{types.B, types.N, types.S}),
 			),
 		},
 	}

ast/env.go

@@ -6,6 +6,7 @@ package ast
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/open-policy-agent/opa/types"
 	"github.com/open-policy-agent/opa/util"
@@ -171,6 +172,11 @@ func (env *TypeEnv) getRefRec(node *typeTreeNode, ref, tail Ref) types.Type {
 	}
 
 	if node.Leaf() {
+		if node.children.Len() > 0 {
+			if child := node.Child(tail[0].Value); child != nil {
+				return env.getRefRec(child, ref, tail[1:])
+			}
+		}
 		return selectRef(node.Value(), tail)
 	}
 
@@ -305,9 +311,9 @@ func (n *typeTreeNode) Put(path Ref, tpe types.Type) {
 }
 
 // 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.
+// If a 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) {
+func (n *typeTreeNode) Insert(path Ref, tpe types.Type, env *TypeEnv) {
 	curr := n
 	for i, term := range path {
 		c, ok := curr.children.Get(term.Value)
@@ -324,7 +330,7 @@ func (n *typeTreeNode) Insert(path Ref, tpe types.Type) {
 				// 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)
+					child.value, err = insertIntoObject(o, path[i+1:], tpe, env)
 					if err != nil {
 						panic(fmt.Errorf("unreachable, insertIntoObject: %w", err))
 					}
@@ -342,7 +348,7 @@ func (n *typeTreeNode) Insert(path Ref, tpe types.Type) {
 		leafs := curr.Leafs()
 		for p, t := range leafs {
 			var err error
-			curr.value, err = insertIntoObject(curr.value.(*types.Object), *p, t)
+			curr.value, err = insertIntoObject(curr.value.(*types.Object), *p, t, env)
 			if err != nil {
 				panic(fmt.Errorf("unreachable, insertIntoObject: %w", err))
 			}
@@ -350,47 +356,62 @@ func (n *typeTreeNode) Insert(path Ref, tpe types.Type) {
 	}
 }
 
-func insertIntoObject(o *types.Object, path Ref, tpe types.Type) (*types.Object, error) {
+func (n *typeTreeNode) String() string {
+	b := strings.Builder{}
+
+	if k := n.key; k != nil {
+		b.WriteString(k.String())
+	} else {
+		b.WriteString("-")
+	}
+
+	if v := n.value; v != nil {
+		b.WriteString(": ")
+		b.WriteString(v.String())
+	}
+
+	n.children.Iter(func(_, v util.T) bool {
+		if child, ok := v.(*typeTreeNode); ok {
+			b.WriteString("\n\t+ ")
+			s := child.String()
+			s = strings.ReplaceAll(s, "\n", "\n\t")
+			b.WriteString(s)
+		}
+		return false
+	})
+
+	return b.String()
+}
+
+func insertIntoObject(o *types.Object, path Ref, tpe types.Type, env *TypeEnv) (*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)
-	}
+	key := env.Get(path[0].Value)
 
 	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
-			}
+		var dynamicProps *types.DynamicProperty
+		if dp := o.DynamicProperties(); dp != nil {
+			dynamicProps = types.NewDynamicProperty(types.Or(o.DynamicProperties().Key, key), types.Or(o.DynamicProperties().Value, tpe))
+		} else {
+			dynamicProps = types.NewDynamicProperty(key, tpe)
 		}
-		staticProps := append(o.StaticProperties(), types.NewStaticProperty(key, tpe))
-		return types.NewObject(staticProps, o.DynamicProperties()), nil
+		return types.NewObject(o.StaticProperties(), dynamicProps), 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)
+	child, err := insertIntoObject(types.NewObject(nil, nil), path[1:], tpe, env)
 	if err != nil {
 		return nil, err
 	}
-	staticProps := append(o.StaticProperties(), types.NewStaticProperty(key, child))
-	return types.NewObject(staticProps, o.DynamicProperties()), nil
+
+	var dynamicProps *types.DynamicProperty
+	if dp := o.DynamicProperties(); dp != nil {
+		dynamicProps = types.NewDynamicProperty(types.Or(o.DynamicProperties().Key, key), types.Or(o.DynamicProperties().Value, child))
+	} else {
+		dynamicProps = types.NewDynamicProperty(key, child)
+	}
+	return types.NewObject(o.StaticProperties(), dynamicProps), nil
 }
 
 func (n *typeTreeNode) Leafs() map[*Ref]types.Type {

ast/env_test.go

@@ -5,7 +5,6 @@
 package ast
 
 import (
-	"encoding/json"
 	"testing"
 
 	"github.com/open-policy-agent/opa/types"
@@ -25,8 +24,8 @@ func TestInsertIntoObject(t *testing.T) {
 			path: Ref{NewTerm(String("a"))},
 			tpe:  types.S,
 			expected: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.S)},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
 		},
 		{
 			note: "empty path",
@@ -49,9 +48,8 @@ func TestInsertIntoObject(t *testing.T) {
 			expected: types.NewObject(
 				[]*types.StaticProperty{
 					types.NewStaticProperty("a", types.S),
-					types.NewStaticProperty("b", types.S),
 				},
-				nil),
+				types.NewDynamicProperty(types.S, types.S)),
 		},
 		{
 			note: "number key",
@@ -63,53 +61,63 @@ func TestInsertIntoObject(t *testing.T) {
 			expected: types.NewObject(
 				[]*types.StaticProperty{
 					types.NewStaticProperty("a", types.S),
-					types.NewStaticProperty(json.Number("2"), types.S),
 				},
-				nil),
+				types.NewDynamicProperty(types.N, types.S)),
 		},
 		{
-			note: "same path, same type",
+			note: "other type value inserted",
 			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),
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
 			path: Ref{NewTerm(String("a"))},
 			tpe:  types.B,
 			expected: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.Or(types.S, types.B))},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.Any{types.B, types.S})),
 		},
 		{
-			note: "same path, any type inserted",
+			note: "any type value inserted",
 			obj: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.S)},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
 			path: Ref{NewTerm(String("a"))},
 			tpe:  types.A,
 			expected: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.A)},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.A)),
 		},
 		{
-			note: "same path, any type present",
+			note: "other type key inserted",
 			obj: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.A)},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
+			path: Ref{NewTerm(Number("42"))},
+			tpe:  types.S,
+			expected: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.Any{types.N, types.S}, types.S)),
+		},
+		{
+			note: "other type key and value inserted",
+			obj: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
+			path: Ref{NewTerm(Number("42"))},
+			tpe:  types.B,
+			expected: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.Any{types.N, types.S}, types.Any{types.B, types.S})),
+		},
+		{
+			note: "any type value present, string inserted",
+			obj: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S, types.A)),
 			path: Ref{NewTerm(String("a"))},
 			tpe:  types.S,
 			expected: types.NewObject(
-				[]*types.StaticProperty{types.NewStaticProperty("a", types.A)},
-				nil),
+				nil,
+				types.NewDynamicProperty(types.S, types.A)),
 		},
 		{
 			note: "long path",
@@ -121,97 +129,82 @@ func TestInsertIntoObject(t *testing.T) {
 			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),
+				types.NewDynamicProperty(types.S, // b
+					types.NewObject(nil, types.NewDynamicProperty(types.S, // c
+						types.NewObject(nil, types.NewDynamicProperty(types.S, types.S)))))), // d
 		},
 		{
-			note: "long path, full match",
+			note: "long path, dynamic overlap with different key 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),
+				nil,
+				types.NewDynamicProperty(types.N, types.S)),
 			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),
+				nil,
+				types.NewDynamicProperty(types.Any{types.N, types.S}, // b
+					types.Any{types.S,
+						types.NewObject(nil, types.NewDynamicProperty(types.S, // c
+							types.NewObject(nil, types.NewDynamicProperty(types.S, types.S))))})), // d
 		},
 		{
-			note: "long path, full match, different type",
+			note: "long path, dynamic overlap with object",
 			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),
+				nil,
+				types.NewDynamicProperty(types.S,
+					types.NewObject(nil, types.NewDynamicProperty(types.S, types.N)))),
 			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),
+				nil,
+				types.NewDynamicProperty(types.S, // b
+					types.Any{
+						types.NewObject(nil, types.NewDynamicProperty(types.S, types.N)),
+						types.NewObject(nil, types.NewDynamicProperty(types.S, // c
+							types.NewObject(nil, types.NewDynamicProperty(types.S, types.S)))), // d
+					})),
+		},
+		{
+			note: "long path, dynamic overlap with object (2)",
+			obj: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S,
+					types.NewObject(nil, types.NewDynamicProperty(types.S,
+						types.NewObject(nil, types.NewDynamicProperty(types.S, types.N)))))),
+			path: Ref{NewTerm(String("b")), NewTerm(String("c")), NewTerm(String("d"))},
+			tpe:  types.S,
+			expected: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S,
+					types.Any{ // Objects aren't merged, as that would become very complicated if they contain static components
+						types.NewObject(nil, types.NewDynamicProperty(types.S,
+							types.NewObject(nil, types.NewDynamicProperty(types.S, types.N)))),
+						types.NewObject(nil, types.NewDynamicProperty(types.S,
+							types.NewObject(nil, types.NewDynamicProperty(types.S, types.S)))),
+					})),
+		},
+		{
+			note: "long path, dynamic overlap with different value type",
+			obj: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S, types.S)),
+			path: Ref{NewTerm(String("b")), NewTerm(String("c")), NewTerm(String("d"))},
+			tpe:  types.S,
+			expected: types.NewObject(
+				nil,
+				types.NewDynamicProperty(types.S, // b
+					types.Any{types.S,
+						types.NewObject(nil, types.NewDynamicProperty(types.S, // c
+							types.NewObject(nil, types.NewDynamicProperty(types.S, types.S))))})), // d
 		},
 	}
 
+	env := TypeEnv{}
+
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
-			result, err := insertIntoObject(tc.obj, tc.path, tc.tpe)
+			result, err := insertIntoObject(tc.obj, tc.path, tc.tpe, &env)
 			if err != nil {
 				t.Fatalf("Unexpected error: %v", err)
 			}
@@ -223,6 +216,7 @@ func TestInsertIntoObject(t *testing.T) {
 }
 
 func TestTypeTreeInsert(t *testing.T) {
+	env := TypeEnv{}
 	n := newTypeTree()
 
 	abcRef := Ref{NewTerm(String("a")), NewTerm(String("b")), NewTerm(String("c"))}
@@ -242,17 +236,14 @@ func TestTypeTreeInsert(t *testing.T) {
 	// 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}))
+	n.Insert(abRef, types.NewObject(nil, &types.DynamicProperty{Key: types.N, Value: types.S}), &env)
 
 	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},
+		nil,
+		types.NewDynamicProperty(
+			types.Any{types.N, types.S},
+			types.Any{types.B, types.S, types.NewObject(nil, types.NewDynamicProperty(types.S, types.N))}),
 	)
 	if types.Compare(actual, expected) != 0 {
 		t.Fatalf("Expected %v but got %v", expected, actual)
@@ -260,8 +251,8 @@ func TestTypeTreeInsert(t *testing.T) {
 
 	// 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)
+	abfRef := Ref{NewTerm(String("a")), NewTerm(String("b")), NewTerm(Boolean(true))}
+	n.Insert(abfRef, types.S, &env)
 
 	actual = n.Get(abfRef)
 	if types.Compare(actual, types.S) != 0 {
@@ -270,14 +261,10 @@ func TestTypeTreeInsert(t *testing.T) {
 
 	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},
+		nil,
+		types.NewDynamicProperty(
+			types.Any{types.B, types.N, types.S},
+			types.Any{types.B, types.S, types.NewObject(nil, types.NewDynamicProperty(types.S, types.N))}),
 	)
 	if types.Compare(actual, expected) != 0 {
 		t.Fatalf("Expected %v but got %v", expected, actual)

test/cases/testdata/type/test-regressions.yaml

@@ -29,4 +29,29 @@ cases:
       }
   query: data.test.q = x
   want_result:
-    - x: ["bar"]
+    - x: ["bar"]
+- note: regression/dynamic object to static object comparison (https://github.com/open-policy-agent/opa/issues/6138)
+  modules:
+    - |
+      package test
+      
+      obj[k] := v {
+        v := ["a", "b", "c"][k]
+      }
+
+      obj.foo := "bar"
+      
+      obj.baz := true { false }
+      
+      compare {
+        # Comparison with static object that doesn't contain "optional" key.
+        obj == {
+          0: "a",
+          1: "b",
+          2: "c",
+          "foo": "bar"
+        }
+      }
+  query: data.test.compare = x
+  want_result:
+    - x: true