refactor: use generics for ADT

ast/check.go

@@ -10,7 +10,6 @@ import (
 	"strings"
 
 	"github.com/open-policy-agent/opa/types"
-	"github.com/open-policy-agent/opa/util"
 )
 
 type varRewriter func(Ref) Ref
@@ -156,10 +155,10 @@ func (tc *typeChecker) CheckBody(env *TypeEnv, body Body) (*TypeEnv, Errors) {
 // CheckTypes runs type checking on the rules returns a TypeEnv if no errors
 // are found. The resulting TypeEnv wraps the provided one. The resulting
 // TypeEnv will be able to resolve types of refs that refer to rules.
-func (tc *typeChecker) CheckTypes(env *TypeEnv, sorted []util.T, as *AnnotationSet) (*TypeEnv, Errors) {
+func (tc *typeChecker) CheckTypes(env *TypeEnv, sorted []*Rule, as *AnnotationSet) (*TypeEnv, Errors) {
 	env = tc.newEnv(env)
 	for _, s := range sorted {
-		tc.checkRule(env, as, s.(*Rule))
+		tc.checkRule(env, as, s)
 	}
 	tc.errs.Sort()
 	return env, tc.errs
@@ -744,8 +743,8 @@ func (rc *refChecker) checkRef(curr *TypeEnv, node *typeTreeNode, ref Ref, idx i
 
 		case RootDocumentNames.Contains(ref[0]):
 			if idx != 0 {
-				node.Children().Iter(func(_, child util.T) bool {
-					_ = rc.checkRef(curr, child.(*typeTreeNode), ref, idx+1) // ignore error
+				node.Children().Iter(func(_ Value, child *typeTreeNode) bool {
+					_ = rc.checkRef(curr, child, ref, idx+1) // ignore error
 					return false
 				})
 				return nil
@@ -1090,8 +1089,8 @@ func newArgError(loc *Location, builtinName Ref, msg string, have []types.Type,
 }
 
 func getOneOfForNode(node *typeTreeNode) (result []Value) {
-	node.Children().Iter(func(k, _ util.T) bool {
-		result = append(result, k.(Value))
+	node.Children().Iter(func(k Value, v *typeTreeNode) bool {
+		result = append(result, k)
 		return false
 	})
 

ast/check_test.go

@@ -661,7 +661,7 @@ func TestCheckInferenceRules(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
-			var elems []util.T
+			var elems []*Rule
 
 			// Convert test rules into rule slice for call.
 			for i := range tc.rules {
@@ -732,7 +732,7 @@ func TestCheckInferenceOverlapWithRules(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
-			var elems []util.T
+			var elems []*Rule
 
 			// Convert test rules into rule slice for "warmup" call.
 			for i := range tc.rules {
@@ -1325,7 +1325,7 @@ func TestFunctionTypeInferenceUnappliedWithObjectVarKey(t *testing.T) {
 		f(x) = y { y = {x: 1} }
 	`)
 
-	elems := []util.T{
+	elems := []*Rule{
 		module.Rules[0],
 	}
 
@@ -1513,8 +1513,8 @@ func TestCheckErrorOrdering(t *testing.T) {
 		p { data.test.q = 2 }  # type error: bool = number
 	`)
 
-	input := make([]util.T, len(mod.Rules))
-	inputReversed := make([]util.T, len(mod.Rules))
+	input := make([]*Rule, len(mod.Rules))
+	inputReversed := make([]*Rule, len(mod.Rules))
 
 	for i := range input {
 		input[i] = mod.Rules[i]
@@ -1562,7 +1562,7 @@ func newTestEnv(rs []string) *TypeEnv {
 
 	// We preallocate enough for at least the base rules.
 	// Else cases will cause reallocs, but that's okay.
-	elems := make([]util.T, 0, len(rs))
+	elems := make([]*Rule, 0, len(rs))
 
 	for i := range rs {
 		rule := MustParseRule(rs[i])
@@ -2305,7 +2305,7 @@ p { input = "foo" }`}},
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
 			var modules []*Module
-			var elems []util.T
+			var elems []*Rule
 
 			for i, module := range tc.modules {
 				mod, err := ParseModuleWithOpts(fmt.Sprintf("test%d.rego", i+1), module, ParserOptions{