Improve token representation in AST

Previously the AST did not hold the token's source end offset in many
cases. This change adds token end offset by offsetting the end by the text
length from the token start.

common/ast/BUILD.bazel

@@ -28,6 +28,7 @@ go_test(
     name = "go_default_test",
     srcs = [
         "ast_test.go",
+        "ast_export_test.go",
         "conversion_test.go",
         "expr_test.go",
         "navigable_test.go",
@@ -50,5 +51,6 @@ go_test(
         "@org_golang_google_genproto_googleapis_api//expr/v1alpha1:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
         "@org_golang_google_protobuf//encoding/prototext:go_default_library",
+        "@com_github_google_go_cmp//cmp:go_default_library",
     ],
 )

common/ast/ast_export_test.go

@@ -0,0 +1,27 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+
+package ast
+
+type (
+	TestExpr                  = expr
+	TestEntryExpr             = entryExpr
+	TestBaseCallExpr          = baseCallExpr
+	TestBaseSelectExpr        = baseSelectExpr
+	TestBaseListExpr          = baseListExpr
+	TestBaseMapExpr           = baseMapExpr
+	TestBaseMapEntry          = baseMapEntry
+	TestBaseStructExpr        = baseStructExpr
+	TestBaseStructField       = baseStructField
+	TestBaseComprehensionExpr = baseComprehensionExpr
+)

common/ast/ast_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/google/cel-go/common/ast"
 	"github.com/google/cel-go/common/overloads"
 	"github.com/google/cel-go/common/types"
+	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 )
@@ -72,8 +73,30 @@ func TestASTCopy(t *testing.T) {
 		if err != nil {
 			t.Errorf("ast.ToAST() failed: %v", err)
 		}
-		if !reflect.DeepEqual(checked, checkedRoundtrip) {
-			t.Errorf("Roundtrip got %v, wanted %v", checkedRoundtrip, checked)
+		opts := []cmp.Option{
+			cmp.AllowUnexported(
+				ast.AST{},
+				ast.SourceInfo{},
+				types.Type{},
+				ast.TestExpr{},
+				ast.TestEntryExpr{},
+				ast.TestBaseSelectExpr{},
+				ast.TestBaseCallExpr{},
+				ast.TestBaseListExpr{},
+				ast.TestBaseMapExpr{},
+				ast.TestBaseMapEntry{},
+				ast.TestBaseStructExpr{},
+				ast.TestBaseStructField{},
+				ast.TestBaseComprehensionExpr{},
+			),
+			cmp.Comparer(func(a, b ast.OffsetRange) bool {
+				// The checked AST may have an offset stop,
+				// but the round-tripped AST will not.
+				return a.Start == b.Start && (a.Start == a.Stop || b.Start == b.Stop)
+			}),
+		}
+		if !cmp.Equal(checked, checkedRoundtrip, opts...) {
+			t.Errorf("Roundtrip got %v, wanted %v\n%s", checkedRoundtrip, checked, cmp.Diff(checked, checkedRoundtrip, opts...))
 		}
 	}
 }

go.mod

@@ -4,6 +4,7 @@ go 1.18
 
 require (
 	github.com/antlr4-go/antlr/v4 v4.13.0
+	github.com/google/go-cmp v0.6.0
 	github.com/stoewer/go-strcase v1.2.0
 	golang.org/x/text v0.9.0
 	google.golang.org/genproto/googleapis/api v0.0.0-20230803162519-f966b187b2e5

go.sum

@@ -2,7 +2,8 @@ github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8
 github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stoewer/go-strcase v1.2.0 h1:Z2iHWqGXH00XYgqDmNgQbIBxf3wrNq0F3feEy0ainaU=

parser/helper.go

@@ -140,15 +140,12 @@ func (p *parserHelper) id(ctx any) int64 {
 	var offset ast.OffsetRange
 	switch c := ctx.(type) {
 	case antlr.ParserRuleContext:
-		start, stop := c.GetStart(), c.GetStop()
-		if stop == nil {
-			stop = start
-		}
+		start := c.GetStart()
 		offset.Start = p.sourceInfo.ComputeOffset(int32(start.GetLine()), int32(start.GetColumn()))
-		offset.Stop = p.sourceInfo.ComputeOffset(int32(stop.GetLine()), int32(stop.GetColumn()))
+		offset.Stop = offset.Start + int32(len(c.GetText()))
 	case antlr.Token:
 		offset.Start = p.sourceInfo.ComputeOffset(int32(c.GetLine()), int32(c.GetColumn()))
-		offset.Stop = offset.Start
+		offset.Stop = offset.Start + int32(len(c.GetText()))
 	case common.Location:
 		offset.Start = p.sourceInfo.ComputeOffset(int32(c.Line()), int32(c.Column()))
 		offset.Stop = offset.Start