chore(ast): Update WalkFn signature to return bool (#16989)

The boolean return value indicates whether walking/traversing should be continued with child nodes.
This PR does not change any existing traversal behaviour.


Signed-off-by: Christian Haudum <christian.haudum@gmail.com>

pkg/dataobj/querier/store.go

@@ -658,19 +658,17 @@ func buildLogsPredicateFromSampleExpr(expr syntax.SampleExpr) (dataobj.LogsPredi
 		predicate dataobj.LogsPredicate
 		skip      bool
 	)
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch e := e.(type) {
 		case *syntax.BinOpExpr:
 			// we might not encounter BinOpExpr at this point since the lhs and rhs are evaluated separately?
 			skip = true
-			return
 		case *syntax.RangeAggregationExpr:
-			if skip {
-				return
+			if !skip {
+				predicate, e.Left.Left = buildLogsPredicateFromPipeline(e.Left.Left)
 			}
-
-			predicate, e.Left.Left = buildLogsPredicateFromPipeline(e.Left.Left)
 		}
+		return true
 	})
 
 	return predicate, expr

pkg/engine/engine.go

@@ -9,7 +9,7 @@ func canExecuteWithNewEngine(expr syntax.Expr) bool {
 		return false
 	case syntax.LogSelectorExpr:
 		ret := true
-		expr.Walk(func(e syntax.Expr) {
+		expr.Walk(func(e syntax.Expr) bool {
 			switch e.(type) {
 			case *syntax.LineParserExpr, *syntax.LogfmtParserExpr, *syntax.LogfmtExpressionParserExpr, *syntax.JSONExpressionParserExpr:
 				ret = false
@@ -18,6 +18,7 @@ func canExecuteWithNewEngine(expr syntax.Expr) bool {
 			case *syntax.KeepLabelsExpr, *syntax.DropLabelsExpr:
 				ret = false
 			}
+			return true
 		})
 		return ret
 	}

pkg/logql/downstream.go

@@ -140,7 +140,9 @@ func (d DownstreamLogSelectorExpr) Pretty(level int) string {
 }
 
 func (d DownstreamSampleExpr) Walk(f syntax.WalkFn) {
-	f(d)
+	if !f(d) {
+		return
+	}
 	if d.SampleExpr != nil {
 		d.SampleExpr.Walk(f)
 	}
@@ -177,7 +179,9 @@ func (c *ConcatSampleExpr) string(maxDepth int) string {
 }
 
 func (c *ConcatSampleExpr) Walk(f syntax.WalkFn) {
-	f(c)
+	if !f(c) {
+		return
+	}
 	if c.SampleExpr != nil {
 		c.SampleExpr.Walk(f)
 	}
@@ -280,7 +284,9 @@ func (e QuantileSketchEvalExpr) String() string {
 }
 
 func (e *QuantileSketchEvalExpr) Walk(f syntax.WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}
@@ -311,7 +317,9 @@ func (e QuantileSketchMergeExpr) String() string {
 }
 
 func (e *QuantileSketchMergeExpr) Walk(f syntax.WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}
@@ -343,7 +351,9 @@ func (e MergeFirstOverTimeExpr) String() string {
 }
 
 func (e *MergeFirstOverTimeExpr) Walk(f syntax.WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}
@@ -375,7 +385,9 @@ func (e MergeLastOverTimeExpr) String() string {
 }
 
 func (e *MergeLastOverTimeExpr) Walk(f syntax.WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}
@@ -406,7 +418,9 @@ func (e CountMinSketchEvalExpr) String() string {
 }
 
 func (e *CountMinSketchEvalExpr) Walk(f syntax.WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}

pkg/logql/engine.go

@@ -515,13 +515,14 @@ func (q *query) JoinSampleVector(next bool, r StepResult, stepEvaluator StepEval
 
 func (q *query) checkIntervalLimit(expr syntax.SampleExpr, limit time.Duration) error {
 	var err error
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch e := e.(type) {
 		case *syntax.LogRangeExpr:
 			if e.Interval > limit {
 				err = fmt.Errorf("%w: [%s] > [%s]", logqlmodel.ErrIntervalLimit, model.Duration(e.Interval), model.Duration(limit))
 			}
 		}
+		return true
 	})
 	return err
 }

pkg/logql/evaluator.go

@@ -243,13 +243,13 @@ func Sortable(q Params) (bool, error) {
 		return false, nil
 	case syntax.SampleExpr:
 		var sortable bool
-		expr.Walk(func(e syntax.Expr) {
+		expr.Walk(func(e syntax.Expr) bool {
 			if rangeExpr, ok := e.(*syntax.VectorAggregationExpr); ok {
 				if rangeExpr.Operation == syntax.OpTypeSort || rangeExpr.Operation == syntax.OpTypeSortDesc {
 					sortable = true
-					return
 				}
 			}
+			return true
 		})
 		return sortable, nil
 	default:

pkg/logql/optimize.go

@@ -6,12 +6,12 @@ import "github.com/grafana/loki/v3/pkg/logql/syntax"
 func optimizeSampleExpr(expr syntax.SampleExpr) (syntax.SampleExpr, error) {
 	var skip bool
 	// we skip sharding AST for now, it's not easy to clone them since they are not part of the language.
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch e.(type) {
 		case *ConcatSampleExpr, DownstreamSampleExpr, *QuantileSketchEvalExpr, *QuantileSketchMergeExpr, *MergeFirstOverTimeExpr, *MergeLastOverTimeExpr:
 			skip = true
-			return
 		}
+		return true
 	})
 	if skip {
 		return expr, nil
@@ -26,19 +26,19 @@ func optimizeSampleExpr(expr syntax.SampleExpr) (syntax.SampleExpr, error) {
 
 // removeLineformat removes unnecessary line_format within a SampleExpr.
 func removeLineformat(expr syntax.SampleExpr) {
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		rangeExpr, ok := e.(*syntax.RangeAggregationExpr)
 		if !ok {
-			return
+			return true
 		}
 		// bytes operation count bytes of the log line so line_format changes the result.
 		if rangeExpr.Operation == syntax.OpRangeTypeBytes ||
 			rangeExpr.Operation == syntax.OpRangeTypeBytesRate {
-			return
+			return true
 		}
 		pipelineExpr, ok := rangeExpr.Left.Left.(*syntax.PipelineExpr)
 		if !ok {
-			return
+			return true
 		}
 		temp := pipelineExpr.MultiStages[:0]
 		for i, s := range pipelineExpr.MultiStages {
@@ -82,5 +82,6 @@ func removeLineformat(expr syntax.SampleExpr) {
 		if len(pipelineExpr.MultiStages) == 0 {
 			rangeExpr.Left.Left = &syntax.MatchersExpr{Mts: rangeExpr.Left.Left.Matchers()}
 		}
+		return true
 	})
 }

pkg/logql/rangemapper.go

@@ -196,11 +196,12 @@ func (m RangeMapper) Map(expr syntax.SampleExpr, vectorAggrPushdown *syntax.Vect
 // Example: expression `count_over_time({app="foo"}[10m])` returns 10m
 func getRangeInterval(expr syntax.SampleExpr) time.Duration {
 	var rangeInterval time.Duration
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch concrete := e.(type) {
 		case *syntax.RangeAggregationExpr:
 			rangeInterval = concrete.Left.Interval
 		}
+		return true
 	})
 	return rangeInterval
 }
@@ -209,7 +210,7 @@ func getRangeInterval(expr syntax.SampleExpr) time.Duration {
 // such as `| json` or `| logfmt`, that would result in an exploding amount of series in downstream queries.
 func hasLabelExtractionStage(expr syntax.SampleExpr) bool {
 	found := false
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch concrete := e.(type) {
 		case *syntax.LogfmtParserExpr:
 			found = true
@@ -220,6 +221,7 @@ func hasLabelExtractionStage(expr syntax.SampleExpr) bool {
 				found = true
 			}
 		}
+		return true
 	})
 	return found
 }
@@ -297,7 +299,7 @@ func (m RangeMapper) vectorAggrWithRangeDownstreams(expr *syntax.RangeAggregatio
 // Returns the updated downstream ConcatSampleExpr.
 func appendDownstream(downstreams *ConcatSampleExpr, expr syntax.SampleExpr, interval time.Duration, offset time.Duration) *ConcatSampleExpr {
 	sampleExpr := syntax.MustClone(expr)
-	sampleExpr.Walk(func(e syntax.Expr) {
+	sampleExpr.Walk(func(e syntax.Expr) bool {
 		switch concrete := e.(type) {
 		case *syntax.RangeAggregationExpr:
 			concrete.Left.Interval = interval
@@ -305,6 +307,7 @@ func appendDownstream(downstreams *ConcatSampleExpr, expr syntax.SampleExpr, int
 				concrete.Left.Offset = offset
 			}
 		}
+		return true
 	})
 	downstreams = &ConcatSampleExpr{
 		DownstreamSampleExpr: DownstreamSampleExpr{
@@ -319,11 +322,12 @@ func getOffsets(expr syntax.SampleExpr) []time.Duration {
 	// Expect to always find at most 1 offset, so preallocate it accordingly
 	offsets := make([]time.Duration, 0, 1)
 
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch concrete := e.(type) {
 		case *syntax.RangeAggregationExpr:
 			offsets = append(offsets, concrete.Left.Offset)
 		}
+		return true
 	})
 	return offsets
 }

pkg/logql/syntax/ast.go

@@ -389,8 +389,9 @@ func (e *PipelineExpr) Shardable(topLevel bool) bool {
 }
 
 func (e *PipelineExpr) Walk(f WalkFn) {
-	f(e)
-
+	if !f(e) {
+		return
+	}
 	if e.Left != nil {
 		e.Left.Walk(f)
 	}
@@ -543,7 +544,9 @@ func newNestedLineFilterExpr(left *LineFilterExpr, right *LineFilterExpr) *LineF
 }
 
 func (e *LineFilterExpr) Walk(f WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.Left != nil {
 		e.Left.Walk(f)
 	}
@@ -1161,7 +1164,9 @@ func (r LogRangeExpr) String() string {
 func (r *LogRangeExpr) Shardable(topLevel bool) bool { return r.Left.Shardable(topLevel) }
 
 func (r *LogRangeExpr) Walk(f WalkFn) {
-	f(r)
+	if !f(r) {
+		return
+	}
 	if r.Left != nil {
 		r.Left.Walk(f)
 	}
@@ -1473,7 +1478,9 @@ func (e *RangeAggregationExpr) Shardable(topLevel bool) bool {
 }
 
 func (e *RangeAggregationExpr) Walk(f WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.Left != nil {
 		e.Left.Walk(f)
 	}
@@ -1683,7 +1690,9 @@ func (e *VectorAggregationExpr) Shardable(topLevel bool) bool {
 }
 
 func (e *VectorAggregationExpr) Walk(f WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.Left != nil {
 		e.Left.Walk(f)
 	}
@@ -1803,7 +1812,9 @@ func (e *BinOpExpr) Shardable(topLevel bool) bool {
 }
 
 func (e *BinOpExpr) Walk(f WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.SampleExpr != nil {
 		e.SampleExpr.Walk(f)
 	}
@@ -2232,7 +2243,9 @@ func (e *LabelReplaceExpr) Shardable(_ bool) bool {
 }
 
 func (e *LabelReplaceExpr) Walk(f WalkFn) {
-	f(e)
+	if !f(e) {
+		return
+	}
 	if e.Left != nil {
 		e.Left.Walk(f)
 	}
@@ -2377,7 +2390,7 @@ func (e *VectorExpr) MatcherGroups() ([]MatcherRange, error) { return nil, e.err
 func (e *VectorExpr) Extractors() ([]log.SampleExtractor, error) { return []log.SampleExtractor{}, nil }
 
 func ReducesLabels(e Expr) (conflict bool) {
-	e.Walk(func(e Expr) {
+	e.Walk(func(e Expr) bool {
 		switch expr := e.(type) {
 		case *RangeAggregationExpr:
 			if groupingReducesLabels(expr.Grouping) {
@@ -2406,9 +2419,8 @@ func ReducesLabels(e Expr) (conflict bool) {
 					break
 				}
 			}
-		default:
-			return
 		}
+		return true
 	})
 	return
 }
@@ -2512,7 +2524,9 @@ func (m *MultiVariantExpr) Shardable(topLevel bool) bool {
 }
 
 func (m *MultiVariantExpr) Walk(f WalkFn) {
-	f(m)
+	if !f(m) {
+		return
+	}
 
 	if m.logRange != nil {
 		m.logRange.Walk(f)

pkg/logql/syntax/test_utils.go

@@ -45,7 +45,7 @@ func removeFastRegexMatcherFromExpr(expr Expr) Expr {
 	if expr == nil {
 		return nil
 	}
-	expr.Walk(func(e Expr) {
+	expr.Walk(func(e Expr) bool {
 		switch typed := e.(type) {
 		case *MatchersExpr:
 			typed.Mts = RemoveFastRegexMatchers(typed.Mts)
@@ -53,7 +53,7 @@ func removeFastRegexMatcherFromExpr(expr Expr) Expr {
 			typed.LabelFilterer = removeFastRegexMatcherFromLabelFilterer(typed.LabelFilterer)
 		case *LogRangeExpr:
 			if typed.Unwrap == nil {
-				return
+				return true
 			}
 			cleaned := make([]log.LabelFilterer, 0, len(typed.Unwrap.PostFilters))
 			for _, filter := range typed.Unwrap.PostFilters {
@@ -65,9 +65,8 @@ func removeFastRegexMatcherFromExpr(expr Expr) Expr {
 			for i, variant := range typed.variants {
 				typed.variants[i] = removeFastRegexMatcherFromExpr(variant).(SampleExpr)
 			}
-		default:
-			return
 		}
+		return true
 	})
 	return expr
 }

pkg/logql/syntax/walk.go

@@ -1,7 +1,10 @@
 package syntax
 
-type WalkFn = func(e Expr)
+// WalkFn is the callback function that gets called whenever a node of the AST is visited.
+// The return value indicates whether the traversal should continue with the child nodes.
+type WalkFn = func(e Expr) bool
 
+// Walkable denotes a node of the AST that can be traversed.
 type Walkable interface {
 	Walk(f WalkFn)
 }

pkg/logql/syntax/walk_test.go

@@ -46,7 +46,7 @@ func Test_Walkable(t *testing.T) {
 			require.Nil(t, err)
 
 			var cnt int
-			expr.Walk(func(_ Expr) { cnt++ })
+			expr.Walk(func(_ Expr) bool { cnt++; return true })
 			require.Equal(t, test.want, cnt)
 		})
 	}
@@ -90,13 +90,14 @@ func Test_AppendMatchers(t *testing.T) {
 			expr, err := ParseExpr(test.expr)
 			require.NoError(t, err)
 
-			expr.Walk(func(e Expr) {
+			expr.Walk(func(e Expr) bool {
 				switch me := e.(type) {
 				case *MatchersExpr:
 					me.AppendMatchers(test.matchers)
 				default:
 					// do nothing
 				}
+				return true
 			})
 			require.Equal(t, test.want, expr.String())
 		})

pkg/querier/multi_tenant_querier.go

@@ -348,11 +348,12 @@ func removeTenantSelector(params logql.SelectSampleParams, tenantIDs []string) (
 // replaceMatchers traverses the passed expression and replaces all matchers.
 func replaceMatchers(expr syntax.Expr, matchers []*labels.Matcher) syntax.Expr {
 	expr, _ = syntax.Clone(expr)
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch concrete := e.(type) {
 		case *syntax.MatchersExpr:
 			concrete.Mts = matchers
 		}
+		return true
 	})
 	return expr
 }

pkg/querier/queryrange/downstreamer.go

@@ -113,7 +113,7 @@ func withoutOffset(query logql.DownstreamQuery) (string, time.Time, time.Time) {
 		newStart = query.Params.Start()
 		newEnd   = query.Params.End()
 	)
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		switch rng := e.(type) {
 		case *syntax.RangeAggregationExpr:
 			off := rng.Left.Offset
@@ -127,6 +127,7 @@ func withoutOffset(query logql.DownstreamQuery) (string, time.Time, time.Time) {
 
 			}
 		}
+		return true
 	})
 	return expr.String(), newStart, newEnd
 }

pkg/querier/queryrange/split_by_interval.go

@@ -265,7 +265,7 @@ func maxRangeVectorAndOffsetDuration(expr syntax.Expr) (time.Duration, time.Dura
 	}
 
 	var maxRVDuration, maxOffset time.Duration
-	expr.Walk(func(e syntax.Expr) {
+	expr.Walk(func(e syntax.Expr) bool {
 		if r, ok := e.(*syntax.LogRangeExpr); ok {
 			if r.Interval > maxRVDuration {
 				maxRVDuration = r.Interval
@@ -274,6 +274,7 @@ func maxRangeVectorAndOffsetDuration(expr syntax.Expr) (time.Duration, time.Dura
 				maxOffset = r.Offset
 			}
 		}
+		return true
 	})
 	return maxRVDuration, maxOffset
 }