Clean up unused source info after checker rewrites the AST.

checker/checker.go

@@ -67,6 +67,15 @@ 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.
+	ids := c.AST.IDs()
+	for id := range c.AST.SourceInfo().OffsetRanges() {
+		if !ids[id] {
+			c.AST.SourceInfo().ClearOffsetRange(id)
+		}
+	}
 	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()
+			missingIDs := []int64{}
+			for id := range cAst.SourceInfo().OffsetRanges() {
+				if !astIDs[id] {
+					missingIDs = append(missingIDs, id)
+				}
+			}
+			if len(missingIDs) > 0 {
+				t.Errorf("SourceInfo has offset range for IDs %v, but no such nodes exists in AST: %s",
+					missingIDs, 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,16 @@ 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
+}
+
 // 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 +543,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{})
+}