General refs in rule heads (#5913)

* Adding support for multiple variables at arbitrary locations in rule refs
* Updating type-checker to handle general ref heads

Fixes: #5993
Fixes: #5994

Signed-off-by: Johan Fylling <[email protected]>

Author: johanfylling

ast/check.go

@@ -231,37 +231,30 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) {
 
 		f := types.NewFunction(args, cpy.Get(rule.Head.Value))
 
-		// Union with existing.
-		exist := env.tree.Get(path)
-		tpe = types.Or(exist, f)
-
+		tpe = f
 	} else {
 		switch rule.Head.RuleKind() {
 		case SingleValue:
 			typeV := cpy.Get(rule.Head.Value)
-			if last := path[len(path)-1]; !last.IsGround() {
-
-				// e.g. store object[string: whatever] at data.p.q.r, not data.p.q.r[x]
+			if !path.IsGround() {
+				// e.g. store object[string: whatever] at data.p.q.r, not data.p.q.r[x] or data.p.q.r[x].y[z]
+				objPath := path.DynamicSuffix()
 				path = path.GroundPrefix()
 
-				typeK := cpy.Get(last)
-				if typeK != nil && typeV != nil {
-					exist := env.tree.Get(path)
-					typeV = types.Or(types.Values(exist), typeV)
-					typeK = types.Or(types.Keys(exist), typeK)
-					tpe = types.NewObject(nil, types.NewDynamicProperty(typeK, typeV))
+				var err error
+				tpe, err = nestedObject(cpy, objPath, typeV)
+				if err != nil {
+					tc.err([]*Error{NewError(TypeErr, rule.Head.Location, err.Error())})
+					tpe = nil
 				}
 			} else {
 				if typeV != nil {
-					exist := env.tree.Get(path)
-					tpe = types.Or(typeV, exist)
+					tpe = typeV
 				}
 			}
 		case MultiValue:
 			typeK := cpy.Get(rule.Head.Key)
 			if typeK != nil {
-				exist := env.tree.Get(path)
-				typeK = types.Or(types.Keys(exist), typeK)
 				tpe = types.NewSet(typeK)
 			}
 		}
@@ -272,6 +265,32 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) {
 	}
 }
 
+// nestedObject creates a nested structure of object types, where each term on path corresponds to a level in the
+// nesting. Each term in the path only contributes to the dynamic portion of its corresponding object.
+func nestedObject(env *TypeEnv, path Ref, tpe types.Type) (types.Type, error) {
+	if len(path) == 0 {
+		return tpe, nil
+	}
+
+	k := path[0]
+	typeV, err := nestedObject(env, path[1:], tpe)
+	if err != nil {
+		return nil, err
+	}
+	if typeV == nil {
+		return nil, nil
+	}
+
+	var dynamicProperty *types.DynamicProperty
+	typeK := env.Get(k)
+	if typeK == nil {
+		return nil, nil
+	}
+	dynamicProperty = types.NewDynamicProperty(typeK, typeV)
+
+	return types.NewObject(nil, dynamicProperty), nil
+}
+
 func (tc *typeChecker) checkExpr(env *TypeEnv, expr *Expr) *Error {
 	if err := tc.checkExprWith(env, expr, 0); err != nil {
 		return err

ast/check_test.go

@@ -358,6 +358,16 @@ func TestCheckInferenceRules(t *testing.T) {
 		{`overlap`, `p.q2.a = input.a { true }`},
 		{`overlap`, `p.q2[56] = input.a { true }`},
 	}
+	ruleset3 := [][2]string{
+		{`simple`, `p.q[r][s] = 42 { x = ["a", "b"]; r = x[s] }`},
+		{`mixed`, `p.q[r].s[t] = 42 { x = ["a", "b"]; r = x[t] }`},
+		{`overrides`, `p.q[r] = "foo" { x = ["a", "b"]; r = x[_] }`},
+		{`overrides`, `p.q.r[s] = 42 { x = ["a", "b"]; x[s] }`},
+		{`overrides`, `p.q[r].s = true { x = [true, false]; r = x[_] }`},
+		{`overrides_static`, `p.q[r].a = "foo" { r = "bar"; s = "baz" }`},
+		{`overrides_static`, `p.q[r].b = 42 { r = "bar" }`},
+		{`overrides_static`, `p.q[r].c = true { r = "bar" }`},
+	}
 
 	tests := []struct {
 		note     string
@@ -549,6 +559,104 @@ func TestCheckInferenceRules(t *testing.T) {
 				types.NewDynamicProperty(types.Any{types.N, types.S}, types.Any{types.B, types.N, types.S}),
 			),
 		},
+		{
+			note:  "general ref-rules, only vars in obj-path, complete obj access",
+			rules: ruleset3,
+			ref:   "data.simple.p.q",
+			expected: types.NewObject(
+				[]*types.StaticProperty{},
+				types.NewDynamicProperty(types.S,
+					types.NewObject(
+						[]*types.StaticProperty{},
+						types.NewDynamicProperty(types.N, types.N),
+					),
+				),
+			),
+		},
+		{
+			note:  "general ref-rules, only vars in obj-path, intermediate obj access",
+			rules: ruleset3,
+			ref:   "data.simple.p.q.b",
+			expected: types.NewObject(
+				[]*types.StaticProperty{},
+				types.NewDynamicProperty(types.N, types.N),
+			),
+		},
+		{
+			note:     "general ref-rules, only vars in obj-path, leaf access",
+			rules:    ruleset3,
+			ref:      "data.simple.p.q.b[1]",
+			expected: types.N,
+		},
+		{
+			note:  "general ref-rules, vars and constants in obj-path, complete obj access",
+			rules: ruleset3,
+			ref:   "data.mixed.p.q",
+			expected: types.NewObject(
+				[]*types.StaticProperty{},
+				types.NewDynamicProperty(types.S,
+					types.NewObject(nil,
+						types.NewDynamicProperty(types.S, types.NewObject(nil,
+							types.NewDynamicProperty(types.N, types.N))),
+					),
+				),
+			),
+		},
+		{
+			note:  "general ref-rules, key overrides, complete obj access",
+			rules: ruleset3,
+			ref:   "data.overrides.p.q",
+			expected: types.NewObject(nil, types.NewDynamicProperty(
+				types.Or(types.B, types.S),
+				types.Any{
+					types.S,
+					types.NewObject(nil, types.NewDynamicProperty(
+						types.Any{types.N, types.S},
+						types.Any{types.B, types.N})),
+				},
+			),
+			),
+		},
+		{
+			note:  "general ref-rules, multiple static key overrides, complete obj access",
+			rules: ruleset3,
+			ref:   "data.overrides_static.p.q",
+			expected: types.NewObject(
+				[]*types.StaticProperty{},
+				types.NewDynamicProperty(types.S,
+					types.NewObject(
+						nil,
+						types.NewDynamicProperty(types.S, types.Any{types.B, types.N, types.S}),
+					),
+				),
+			),
+		},
+		{
+			note:  "general ref-rules, multiple static key overrides, intermediate obj access",
+			rules: ruleset3,
+			ref:   "data.overrides_static.p.q.foo",
+			expected: types.NewObject(nil,
+				types.NewDynamicProperty(types.S, types.Any{types.B, types.N, types.S}),
+			),
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (a)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.foo.a",
+			expected: types.Any{types.B, types.N, types.S}, // Dynamically build object types don't have static properties, so even though we "know" the 'a' key has a string value, we've lost this information.
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (b)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.bar.b",
+			expected: types.Any{types.B, types.N, types.S},
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (c)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.baz.c",
+			expected: types.Any{types.B, types.N, types.S},
+		},
 	}
 
 	for _, tc := range tests {