Clean up unused source info after checker rewrites the AST. (#1258)

* Clean up unused source info after checker rewrites the AST.

* Move ClearUnusedIDs() to AST.

Also add a check for unused IDs in parser_test.go

* Add comment for ClearUnusedIDs.

Author: abenea

checker/checker.go

@@ -67,6 +67,10 @@ func Check(parsed *ast.AST, source common.Source, env *Env) (*ast.AST, *common.E
 	for id, t := range c.TypeMap() {
 		c.SetType(id, substitute(c.mappings, t, true))
 	}
+	// Remove source info for IDs without a corresponding AST node. This can happen because
+	// check() deletes some nodes while rewriting the AST. For example the Select operand is
+	// deleted when a variable reference is replaced with a Ident expression.
+	c.AST.ClearUnusedIDs()
 	return c.AST, errs
 }
 

checker/checker_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/google/cel-go/common"
 	"github.com/google/cel-go/common/ast"
 	"github.com/google/cel-go/common/containers"
+	"github.com/google/cel-go/common/debug"
 	"github.com/google/cel-go/common/decls"
 	"github.com/google/cel-go/common/stdlib"
 	"github.com/google/cel-go/common/types"
@@ -2553,6 +2554,18 @@ func TestCheck(t *testing.T) {
 				t.Errorf("Expected error not thrown: %s", tc.err)
 			}
 
+			astIDs := cAst.IDs()
+			unusedIDs := []int64{}
+			for id := range cAst.SourceInfo().OffsetRanges() {
+				if !astIDs[id] {
+					unusedIDs = append(unusedIDs, id)
+				}
+			}
+			if len(unusedIDs) > 0 {
+				t.Errorf("SourceInfo has offset range for IDs %v, but no such nodes exists in AST: %s",
+					unusedIDs, debug.ToDebugStringWithIDs(cAst.Expr()))
+			}
+
 			actual := cAst.GetType(pAst.Expr().ID())
 			if tc.err == "" {
 				if actual == nil || !actual.IsEquivalentType(tc.outType) {

common/ast/ast.go

@@ -160,6 +160,26 @@ func MaxID(a *AST) int64 {
 	return visitor.maxID + 1
 }
 
+// IDs returns the set of AST node IDs, including macro calls.
+func (a *AST) IDs() map[int64]bool {
+	visitor := make(idVisitor)
+	PostOrderVisit(a.Expr(), visitor)
+	for _, call := range a.SourceInfo().MacroCalls() {
+		PostOrderVisit(call, visitor)
+	}
+	return visitor
+}
+
+// ClearUnusedIDs removes IDs not used in the AST or macro calls from SourceInfo.
+func (a *AST) ClearUnusedIDs() {
+	ids := a.IDs()
+	for id := range a.SourceInfo().OffsetRanges() {
+		if !ids[id] {
+			a.SourceInfo().ClearOffsetRange(id)
+		}
+	}
+}
+
 // Heights computes the heights of all AST expressions and returns a map from expression id to height.
 func Heights(a *AST) map[int64]int {
 	visitor := make(heightVisitor)
@@ -533,3 +553,13 @@ func (hv heightVisitor) maxEntryHeight(entries ...EntryExpr) int {
 	}
 	return max
 }
+
+type idVisitor map[int64]bool
+
+func (v idVisitor) VisitExpr(e Expr) {
+	v[e.ID()] = true
+}
+
+func (v idVisitor) VisitEntryExpr(e EntryExpr) {
+	v[e.ID()] = true
+}

common/debug/debug.go

@@ -312,3 +312,18 @@ func (w *debugWriter) removeIndent() {
 func (w *debugWriter) String() string {
 	return w.buffer.String()
 }
+
+type idAdorner struct{}
+
+func (a *idAdorner) GetMetadata(elem any) string {
+	e, isExpr := elem.(ast.Expr)
+	if !isExpr {
+		return ""
+	}
+	return fmt.Sprintf("@id:%d ", e.ID())
+}
+
+// ToDebugStringWithIDs returns a string representation with AST node IDs.
+func ToDebugStringWithIDs(e ast.Expr) string {
+	return ToAdornedDebugString(e, &idAdorner{})
+}

parser/parser_test.go

@@ -2262,6 +2262,18 @@ func TestParse(t *testing.T) {
 					t.Fatal(test.DiffMessage(fmt.Sprintf("Macro Calls - %s", failureDisplayMethod), actualAdornedMacroCalls, tc.M))
 				}
 			}
+
+			astIDs := parsed.IDs()
+			unusedIDs := []int64{}
+			for id := range parsed.SourceInfo().OffsetRanges() {
+				if !astIDs[id] {
+					unusedIDs = append(unusedIDs, id)
+				}
+			}
+			if len(unusedIDs) > 0 {
+				t.Errorf("SourceInfo has offset range for IDs %v, but no such nodes exists in AST: %s",
+					unusedIDs, debug.ToDebugStringWithIDs(parsed.Expr()))
+			}
 		})
 	}
 }