Fix issues discovered through fuzz testing

ast/check.go

@@ -669,10 +669,6 @@ func (rc *refChecker) checkRefLeaf(tpe types.Type, ref Ref, idx int) *Error {
 		}
 
 	case *Array, Object, Set:
-		// Composite references operands may only be used with a set.
-		if !unifies(tpe, types.NewSet(types.A)) {
-			return newRefErrInvalid(ref[0].Location, rc.varRewriter(ref), idx, tpe, types.NewSet(types.A), nil)
-		}
 		if !unify1(rc.env, head, keys, false) {
 			return newRefErrInvalid(ref[0].Location, rc.varRewriter(ref), idx, rc.env.Get(head), keys, nil)
 		}

ast/check_test.go

@@ -110,6 +110,16 @@ func TestCheckInference(t *testing.T) {
 				nil,
 			),
 		}},
+		{"object-composite-ref-operand", `x = {{}: 1}; x[{}] = y`, map[Var]types.Type{
+			Var("x"): types.NewObject(
+				[]*types.StaticProperty{types.NewStaticProperty(
+					map[string]interface{}{},
+					types.N,
+				)},
+				nil,
+			),
+			Var("y"): types.N,
+		}},
 		{"sets", `x = {1, 2}; y = {{"foo", 1}, x}`, map[Var]types.Type{
 			Var("x"): types.NewSet(types.N),
 			Var("y"): types.NewSet(
@@ -795,12 +805,12 @@ func TestCheckRefErrInvalid(t *testing.T) {
 			oneOf: []Value{String("p"), String("q")},
 		},
 		{
-			note:  "composite ref into non-set",
+			note:  "composite ref operand",
 			query: `data.test.q[[1, 2]]`,
 			ref:   "data.test.q[[1, 2]]",
 			pos:   3,
-			have:  types.NewObject([]*types.StaticProperty{types.NewStaticProperty("bar", types.N), types.NewStaticProperty("foo", types.N)}, nil),
-			want:  types.NewSet(types.A),
+			have:  types.NewArray([]types.Type{types.N, types.N}, nil),
+			want:  types.S,
 		},
 		{
 			note:  "composite ref type error 1",
@@ -818,6 +828,14 @@ func TestCheckRefErrInvalid(t *testing.T) {
 			have:  types.NewObject([]*types.StaticProperty{types.NewStaticProperty("a", types.S)}, nil),
 			want:  types.NewObject([]*types.StaticProperty{types.NewStaticProperty("a", types.N)}, nil),
 		},
+		{
+			note:  "composite ref type error 3 - array",
+			query: `a = [1,2,3]; a[{}] = b`,
+			ref:   `a[{}]`,
+			pos:   1,
+			have:  types.NewObject(nil, types.NewDynamicProperty(types.A, types.A)),
+			want:  types.N,
+		},
 	}
 
 	for _, tc := range tests {

ast/compile.go

@@ -3033,6 +3033,13 @@ func rewriteDynamicsComprehensionBody(original *Expr, f *equalityFactory, body B
 }
 
 func rewriteExprTermsInHead(gen *localVarGenerator, rule *Rule) {
+	for i := range rule.Head.Args {
+		support, output := expandExprTerm(gen, rule.Head.Args[i])
+		for j := range support {
+			rule.Body.Append(support[j])
+		}
+		rule.Head.Args[i] = output
+	}
 	if rule.Head.Key != nil {
 		support, output := expandExprTerm(gen, rule.Head.Key)
 		for i := range support {

ast/compile_test.go

@@ -1197,6 +1197,17 @@ func TestCompilerRewriteExprTerms(t *testing.T) {
 			`,
 			expected: Errors{&Error{Message: "rule arguments cannot contain calls"}},
 		},
+		{
+			note: "indirect ref in args",
+			module: `
+				package test
+
+				f([1][0]) { true }`,
+			expected: `
+				package test
+
+				f(__local0__[0]) { true; __local0__ = [1] }`,
+		},
 	}
 
 	for _, tc := range cases {

types/types.go

@@ -822,6 +822,15 @@ func TypeOf(x interface{}) Type {
 		return S
 	case json.Number:
 		return N
+	case map[string]interface{}:
+		// The ast.ValueToInterface() function returns ast.Object values as map[string]interface{}
+		// so map[string]interface{} must be handled here because the type checker uses the value
+		// to interface conversion when inferring object types.
+		static := make([]*StaticProperty, 0, len(x))
+		for k, v := range x {
+			static = append(static, NewStaticProperty(k, TypeOf(v)))
+		}
+		return NewObject(static, nil)
 	case map[interface{}]interface{}:
 		static := make([]*StaticProperty, 0, len(x))
 		for k, v := range x {

types/types_test.go

@@ -298,6 +298,22 @@ func TestTypeOf(t *testing.T) {
 	}
 }
 
+func TestTypeOfMapOfString(t *testing.T) {
+	tpe := TypeOf(map[string]interface{}{
+		"foo": "bar",
+		"baz": "qux",
+	})
+
+	exp := NewObject([]*StaticProperty{
+		NewStaticProperty("foo", S),
+		NewStaticProperty("baz", S),
+	}, nil)
+
+	if Compare(exp, tpe) != 0 {
+		t.Fatalf("Expected %v but got: %v", exp, tpe)
+	}
+}
+
 func TestNil(t *testing.T) {
 
 	tpe := NewObject([]*StaticProperty{