gistfile1.txt

diff --git a/internal/planner/planner.go b/internal/planner/planner.go
index eefcd8fb..74fa337b 100644
--- a/internal/planner/planner.go
+++ b/internal/planner/planner.go
@@ -152,10 +152,14 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) {
 	pcurr := p.curr
 	pltarget := p.ltarget
 	plnext := p.lnext
+	ploc := p.loc
 
 	// Reset the variable counter for the function plan.
 	p.lnext = ir.Input
 
+	// Set the location to the rule head.
+	p.loc = rules[0].Head.Loc()
+
 	// Create function definition for rules.
 	fn := &ir.Func{
 		Name: fmt.Sprintf("g%d.%s", p.funcs.gen, path),
@@ -228,6 +232,9 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) {
 		// Unordered rules are treated as a special case of ordered rules.
 		for rule := rules[i]; rule != nil; prev, rule = rule, rule.Else {
 
+			// Update the location for each ordered rule.
+			p.loc = rule.Head.Loc()
+
 			// Setup planner for block.
 			p.lnext = lnext
 			p.vars = newVarstack(map[ast.Var]ir.Local{
@@ -304,6 +311,9 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) {
 	// Default rules execute if the return is undefined.
 	if defaultRule != nil {
 
+		// Set the location for the default rule head.
+		p.loc = defaultRule.Head.Loc()
+
 		fn.Blocks = append(fn.Blocks, &ir.Block{
 			Stmts: []ir.Stmt{
 				&ir.IsUndefinedStmt{Source: fn.Return},
@@ -344,6 +354,7 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) {
 	p.ltarget = pltarget
 	p.vars = pvars
 	p.curr = pcurr
+	p.loc = ploc
 
 	return fn.Name, nil
 }
@@ -450,6 +461,12 @@ func (p *Planner) planQuery(q ast.Body, index int, iter planiter) error {
 		return iter()
 	}
 
+	old := p.loc
+	p.loc = q[index].Loc()
+	defer func() {
+		p.loc = old
+	}()
+
 	return p.planExpr(q[index], func() error {
 		return p.planQuery(q, index+1, iter)
 	})
@@ -457,7 +474,6 @@ func (p *Planner) planQuery(q ast.Body, index int, iter planiter) error {
 
 // TODO(tsandall): improve errors to include location information.
 func (p *Planner) planExpr(e *ast.Expr, iter planiter) error {
-	p.setLocation(e.Location)
 
 	if e.Negated {
 		return p.planNot(e, iter)