Arbitrary vars in rule refs

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,39 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) {
 	}
 }
 
+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 staticProperties []*types.StaticProperty
+	var dynamicProperty *types.DynamicProperty
+	if k.IsGround() {
+		key, err := JSON(path[0].Value)
+		if err != nil {
+			return nil, err
+		}
+		staticProperties = append(staticProperties, types.NewStaticProperty(key, typeV))
+	} else {
+		typeK := env.Get(k)
+		if typeK == nil {
+			return nil, nil
+		}
+		dynamicProperty = types.NewDynamicProperty(typeK, typeV)
+	}
+
+	return types.NewObject(staticProperties, 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.q.a = input.a { true }`},
 		{`overlap`, `p.q[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
@@ -554,6 +564,121 @@ func TestCheckInferenceRules(t *testing.T) {
 				types.NewDynamicProperty(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(
+						[]*types.StaticProperty{types.NewStaticProperty("s",
+							types.NewObject(
+								[]*types.StaticProperty{},
+								types.NewDynamicProperty(types.N, types.N),
+							))},
+						nil,
+					),
+				),
+			),
+		},
+		{
+			note:  "general ref-rules, key overrides, complete obj access",
+			rules: ruleset3,
+			ref:   "data.overrides.p.q",
+			expected: types.NewObject(
+				[]*types.StaticProperty{types.NewStaticProperty("r",
+					types.NewObject(
+						nil,
+						types.NewDynamicProperty(types.N, types.N),
+					))},
+				types.NewDynamicProperty(types.Or(types.B, types.S),
+					types.Or(
+						types.S,
+						types.NewObject(
+							[]*types.StaticProperty{types.NewStaticProperty("s", types.B)},
+							nil,
+						),
+					),
+				),
+			),
+		},
+		{
+			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(
+						[]*types.StaticProperty{
+							types.NewStaticProperty("a", types.S),
+							types.NewStaticProperty("b", types.N),
+							types.NewStaticProperty("c", types.B)},
+						nil,
+					),
+				),
+			),
+		},
+		{
+			note:  "general ref-rules, multiple static key overrides, intermediate obj access",
+			rules: ruleset3,
+			ref:   "data.overrides_static.p.q.foo",
+			expected: types.NewObject(
+				[]*types.StaticProperty{
+					types.NewStaticProperty("a", types.S),
+					types.NewStaticProperty("b", types.N),
+					types.NewStaticProperty("c", types.B)},
+				nil,
+			),
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (a)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.foo.a",
+			expected: types.S,
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (b)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.bar.b",
+			expected: types.N,
+		},
+		{
+			note:     "general ref-rules, multiple static key overrides, leaf access (c)",
+			rules:    ruleset3,
+			ref:      "data.overrides_static.p.q.baz.c",
+			expected: types.B,
+		},
 	}
 
 	for _, tc := range tests {