Populate macro calls

common/source.go

@@ -39,6 +39,10 @@ type Source interface {
 	// and second line, or EOF if there is only one line of source.
 	LineOffsets() []int32
 
+	// Macro calls returns the macro calls map containing the original
+	// expression from a macro replacement, keyed by Id.
+	MacroCalls() map[int64]*exprpb.Expr
+
 	// LocationOffset translates a Location to an offset.
 	// Given the line and column of the Location returns the
 	// Location's character offset in the Source, and a bool
@@ -65,6 +69,7 @@ type sourceImpl struct {
 	description string
 	lineOffsets []int32
 	idOffsets   map[int64]int32
+	macroCalls  map[int64]*exprpb.Expr
 }
 
 var _ runes.Buffer = &sourceImpl{}
@@ -93,6 +98,7 @@ func NewStringSource(contents string, description string) Source {
 		description: description,
 		lineOffsets: offsets,
 		idOffsets:   map[int64]int32{},
+		macroCalls:  map[int64]*exprpb.Expr{},
 	}
 }
 
@@ -103,6 +109,7 @@ func NewInfoSource(info *exprpb.SourceInfo) Source {
 		description: info.GetLocation(),
 		lineOffsets: info.GetLineOffsets(),
 		idOffsets:   info.GetPositions(),
+		macroCalls:  info.GetMacroCalls(),
 	}
 }
 
@@ -121,6 +128,11 @@ func (s *sourceImpl) LineOffsets() []int32 {
 	return s.lineOffsets
 }
 
+// MacroCalls implements the Source interface method.
+func (s *sourceImpl) MacroCalls() map[int64]*exprpb.Expr {
+	return s.macroCalls
+}
+
 // LocationOffset implements the Source interface method.
 func (s *sourceImpl) LocationOffset(location Location) (int32, bool) {
 	if lineOffset, found := s.findLineOffset(location.Line()); found {

parser/helper.go

@@ -41,7 +41,8 @@ func (p *parserHelper) getSourceInfo() *exprpb.SourceInfo {
 	return &exprpb.SourceInfo{
 		Location:    p.source.Description(),
 		Positions:   p.positions,
-		LineOffsets: p.source.LineOffsets()}
+		LineOffsets: p.source.LineOffsets(),
+		MacroCalls:  p.source.MacroCalls()}
 }
 
 func (p *parserHelper) newLiteral(ctx interface{}, value *exprpb.Constant) *exprpb.Expr {
@@ -207,6 +208,62 @@ func (p *parserHelper) getLocation(id int64) common.Location {
 	return location
 }
 
+// buildMacroCallArg iterates the expression and returns a new expression
+// where all macros have been replaced by their IDs in MacroCalls
+func (p *parserHelper) buildMacroCallArg(expr *exprpb.Expr) *exprpb.Expr {
+	resultExpr := &exprpb.Expr{Id: expr.GetId()}
+	if _, found := p.source.MacroCalls()[expr.GetId()]; found {
+		return resultExpr
+	}
+
+	switch expr.ExprKind.(type) {
+	case *exprpb.Expr_CallExpr:
+		resultExpr.ExprKind = &exprpb.Expr_CallExpr{
+			CallExpr: &exprpb.Expr_Call{
+				Function: expr.GetCallExpr().GetFunction(),
+			},
+		}
+		resultExpr.GetCallExpr().Args = make([]*exprpb.Expr, len(expr.GetCallExpr().GetArgs()))
+		// Iterate the AST from `expr` recursively looking for macros. Because we are at most
+		// starting from the top level macro, this recursion is bounded by the size of the AST. This
+		// means that the depth check on the AST during parsing will catch recursion overflows
+		// before we get to here.
+		for index, arg := range expr.GetCallExpr().GetArgs() {
+			resultExpr.GetCallExpr().GetArgs()[index] = p.buildMacroCallArg(arg)
+		}
+		return resultExpr
+	}
+
+	return expr
+}
+
+// addMacroCall adds the macro the the MacroCalls map in source info. If a macro has args/subargs/target
+// that are macros, their ID will be stored instead for later self-lookups.
+func (p *parserHelper) addMacroCall(exprID int64, function string, target *exprpb.Expr, args ...*exprpb.Expr) {
+	expr := &exprpb.Expr{
+		Id: exprID,
+		ExprKind: &exprpb.Expr_CallExpr{
+			CallExpr: &exprpb.Expr_Call{
+				Function: function,
+			},
+		},
+	}
+
+	if target != nil {
+		if _, found := p.source.MacroCalls()[target.GetId()]; found {
+			expr.GetCallExpr().Target = &exprpb.Expr{Id: target.GetId()}
+		} else {
+			expr.GetCallExpr().Target = target
+		}
+	}
+
+	expr.GetCallExpr().Args = make([]*exprpb.Expr, len(args))
+	for index, arg := range args {
+		expr.GetCallExpr().GetArgs()[index] = p.buildMacroCallArg(arg)
+	}
+	p.source.MacroCalls()[exprID] = expr
+}
+
 // balancer performs tree balancing on operators whose arguments are of equal precedence.
 //
 // The purpose of the balancer is to ensure a compact serialization format for the logical &&, ||

parser/macro.go

@@ -285,7 +285,7 @@ func makeExistsOne(eh ExprHelper, target *exprpb.Expr, args []*exprpb.Expr) (*ex
 func makeQuantifier(kind quantifierKind, eh ExprHelper, target *exprpb.Expr, args []*exprpb.Expr) (*exprpb.Expr, *common.Error) {
 	v, found := extractIdent(args[0])
 	if !found {
-		location := eh.OffsetLocation(args[0].Id)
+		location := eh.OffsetLocation(args[0].GetId())
 		return nil, &common.Error{
 			Message:  "argument must be a simple name",
 			Location: location}
@@ -373,14 +373,14 @@ func makeFilter(eh ExprHelper, target *exprpb.Expr, args []*exprpb.Expr) (*exprp
 func extractIdent(e *exprpb.Expr) (string, bool) {
 	switch e.ExprKind.(type) {
 	case *exprpb.Expr_IdentExpr:
-		return e.GetIdentExpr().Name, true
+		return e.GetIdentExpr().GetName(), true
 	}
 	return "", false
 }
 
 func makeHas(eh ExprHelper, target *exprpb.Expr, args []*exprpb.Expr) (*exprpb.Expr, *common.Error) {
 	if s, ok := args[0].ExprKind.(*exprpb.Expr_SelectExpr); ok {
-		return eh.PresenceTest(s.SelectExpr.Operand, s.SelectExpr.Field), nil
+		return eh.PresenceTest(s.SelectExpr.GetOperand(), s.SelectExpr.GetField()), nil
 	}
 	return nil, &common.Error{Message: "invalid argument to has() macro"}
 }

parser/options.go

@@ -22,6 +22,7 @@ type options struct {
 	errorRecoveryLimit               int
 	expressionSizeCodePointLimit     int
 	macros                           map[string]Macro
+	populateMacroCalls               bool
 }
 
 // Option configures the behavior of the parser.
@@ -92,3 +93,12 @@ func Macros(macros ...Macro) Option {
 		return nil
 	}
 }
+
+// PopulateMacroCalls ensures that the original call signatures replaced by expanded macros
+// are preserved in the `SourceInfo` of parse result.
+func PopulateMacroCalls(populateMacroCalls bool) Option {
+	return func(opts *options) error {
+		opts.populateMacroCalls = populateMacroCalls
+		return nil
+	}
+}

parser/parser.go

@@ -66,6 +66,7 @@ func NewParser(opts ...Option) (*Parser, error) {
 	if p.expressionSizeCodePointLimit == -1 {
 		p.expressionSizeCodePointLimit = int((^uint(0)) >> 1)
 	}
+	// Bool is false by default, so populateMacroCalls will be false by default
 	return p, nil
 }
 
@@ -90,6 +91,7 @@ func (p *Parser) Parse(source common.Source) (*exprpb.ParsedExpr, *common.Errors
 		maxRecursionDepth:                p.maxRecursionDepth,
 		errorRecoveryLimit:               p.errorRecoveryLimit,
 		errorRecoveryLookaheadTokenLimit: p.errorRecoveryTokenLookaheadLimit,
+		populateMacroCalls:               p.populateMacroCalls,
 	}
 	buf, ok := source.(runes.Buffer)
 	if !ok {
@@ -278,6 +280,7 @@ type parser struct {
 	maxRecursionDepth                int
 	errorRecoveryLimit               int
 	errorRecoveryLookaheadTokenLimit int
+	populateMacroCalls               bool
 }
 
 var (
@@ -804,15 +807,15 @@ func (p *parser) extractQualifiedName(e *exprpb.Expr) (string, bool) {
 	}
 	switch e.ExprKind.(type) {
 	case *exprpb.Expr_IdentExpr:
-		return e.GetIdentExpr().Name, true
+		return e.GetIdentExpr().GetName(), true
 	case *exprpb.Expr_SelectExpr:
 		s := e.GetSelectExpr()
 		if prefix, found := p.extractQualifiedName(s.Operand); found {
 			return prefix + "." + s.Field, true
 		}
 	}
 	// TODO: Add a method to Source to get location from character offset.
-	location := p.helper.getLocation(e.Id)
+	location := p.helper.getLocation(e.GetId())
 	p.reportError(location, "expected a qualified name")
 	return "", false
 }
@@ -833,7 +836,7 @@ func (p *parser) reportError(ctx interface{}, format string, args ...interface{}
 		location = ctx.(common.Location)
 	case antlr.Token, antlr.ParserRuleContext:
 		err := p.helper.newExpr(ctx)
-		location = p.helper.getLocation(err.Id)
+		location = p.helper.getLocation(err.GetId())
 	}
 	err := p.helper.newExpr(ctx)
 	// Provide arguments to the report error.
@@ -893,5 +896,8 @@ func (p *parser) expandMacro(exprID int64, function string, target *exprpb.Expr,
 		}
 		return p.reportError(p.helper.getLocation(exprID), err.Message), true
 	}
+	if p.populateMacroCalls {
+		p.helper.addMacroCall(expr.GetId(), function, target, args...)
+	}
 	return expr, true
 }