Define visitors.

pkg/logql/syntax/ast.go

@@ -53,11 +53,15 @@ type implicit struct{}
 func (implicit) logQLExpr() {}
 
 // LogSelectorExpr is a LogQL expression filtering and returning logs.
+//
+//sumtype:decl
 type LogSelectorExpr interface {
 	Matchers() []*labels.Matcher
-	LogPipelineExpr
+	Pipeline() (Pipeline, error)
 	HasFilter() bool
 	Expr
+
+	isLogSelectorExpr()
 }
 
 // Type alias for backward compatibility
@@ -66,19 +70,17 @@ type (
 	SampleExtractor = log.SampleExtractor
 )
 
-// LogPipelineExpr is an expression defining a log pipeline.
-type LogPipelineExpr interface {
-	Pipeline() (Pipeline, error)
-	Expr
-}
-
 // StageExpr is an expression defining a single step into a log pipeline
+//
+//sumtype:decl
 type StageExpr interface {
 	Stage() (log.Stage, error)
 	Expr
+
+	isStageExpr()
 }
 
-// MultiStageExpr is multiple stages which implement a PipelineExpr.
+// MultiStageExpr is multiple stages which implements a LogSelectorExpr.
 type MultiStageExpr []StageExpr
 
 func (m MultiStageExpr) Pipeline() (log.Pipeline, error) {
@@ -196,6 +198,8 @@ func newMatcherExpr(matchers []*labels.Matcher) *MatchersExpr {
 	return &MatchersExpr{Mts: matchers}
 }
 
+func (e *MatchersExpr) isLogSelectorExpr() {}
+
 func (e *MatchersExpr) Matchers() []*labels.Matcher {
 	return e.Mts
 }
@@ -242,6 +246,8 @@ func newPipelineExpr(left *MatchersExpr, pipeline MultiStageExpr) LogSelectorExp
 	}
 }
 
+func (e *PipelineExpr) isLogSelectorExpr() {}
+
 func (e *PipelineExpr) Shardable() bool {
 	for _, p := range e.MultiStages {
 		if !p.Shardable() {
@@ -333,6 +339,8 @@ func newNestedLineFilterExpr(left *LineFilterExpr, right *LineFilterExpr) *LineF
 	}
 }
 
+func (*LineFilterExpr) isStageExpr() {}
+
 func (e *LineFilterExpr) Walk(f WalkFn) {
 	f(e)
 	if e.Left == nil {
@@ -471,6 +479,8 @@ func newLogfmtParserExpr(flags []string) *LogfmtParserExpr {
 	return &e
 }
 
+func (*LogfmtParserExpr) isStageExpr() {}
+
 func (e *LogfmtParserExpr) Shardable() bool { return true }
 
 func (e *LogfmtParserExpr) Walk(f WalkFn) { f(e) }
@@ -524,6 +534,8 @@ func newLabelParserExpr(op, param string) *LabelParserExpr {
 	}
 }
 
+func (*LabelParserExpr) isStageExpr() {}
+
 func (e *LabelParserExpr) Shardable() bool { return true }
 
 func (e *LabelParserExpr) Walk(f WalkFn) { f(e) }
@@ -569,6 +581,8 @@ func newLabelFilterExpr(filterer log.LabelFilterer) *LabelFilterExpr {
 	}
 }
 
+func (*LabelFilterExpr) isStageExpr() {}
+
 func (e *LabelFilterExpr) Shardable() bool { return true }
 
 func (e *LabelFilterExpr) Walk(f WalkFn) { f(e) }
@@ -606,6 +620,8 @@ func newDecolorizeExpr() *DecolorizeExpr {
 	return &DecolorizeExpr{}
 }
 
+func (*DecolorizeExpr) isStageExpr() {}
+
 func (e *DecolorizeExpr) Shardable() bool { return true }
 
 func (e *DecolorizeExpr) Stage() (log.Stage, error) {
@@ -625,6 +641,8 @@ func newDropLabelsExpr(dropLabels []log.DropLabel) *DropLabelsExpr {
 	return &DropLabelsExpr{dropLabels: dropLabels}
 }
 
+func (*DropLabelsExpr) isStageExpr() {}
+
 func (e *DropLabelsExpr) Shardable() bool { return true }
 
 func (e *DropLabelsExpr) Stage() (log.Stage, error) {
@@ -663,6 +681,8 @@ func newKeepLabelsExpr(keepLabels []log.KeepLabel) *KeepLabelsExpr {
 	return &KeepLabelsExpr{keepLabels: keepLabels}
 }
 
+func (*KeepLabelsExpr) isStageExpr() {}
+
 func (e *KeepLabelsExpr) Shardable() bool { return true }
 
 func (e *KeepLabelsExpr) Stage() (log.Stage, error) {
@@ -694,6 +714,8 @@ func (e *KeepLabelsExpr) String() string {
 
 func (e *KeepLabelsExpr) Walk(f WalkFn) { f(e) }
 
+func (*LineFmtExpr) isStageExpr() {}
+
 func (e *LineFmtExpr) Shardable() bool { return true }
 
 func (e *LineFmtExpr) Walk(f WalkFn) { f(e) }
@@ -717,6 +739,8 @@ func newLabelFmtExpr(fmts []log.LabelFmt) *LabelFmtExpr {
 	}
 }
 
+func (*LabelFmtExpr) isStageExpr() {}
+
 func (e *LabelFmtExpr) Shardable() bool {
 	// While LabelFmt is shardable in certain cases, it is not always,
 	// but this is left to the shardmapper to determine
@@ -761,6 +785,8 @@ func newJSONExpressionParser(expressions []log.LabelExtractionExpr) *JSONExpress
 	}
 }
 
+func (*JSONExpressionParser) isStageExpr() {}
+
 func (j *JSONExpressionParser) Shardable() bool { return true }
 
 func (j *JSONExpressionParser) Walk(f WalkFn) { f(j) }
@@ -813,6 +839,8 @@ func newLogfmtExpressionParser(expressions []log.LabelExtractionExpr, flags []st
 	return &e
 }
 
+func (*LogfmtExpressionParser) isStageExpr() {}
+
 func (l *LogfmtExpressionParser) Shardable() bool { return true }
 
 func (l *LogfmtExpressionParser) Walk(f WalkFn) { f(l) }
@@ -1101,12 +1129,15 @@ func IsLogicalBinOp(op string) bool {
 }
 
 // SampleExpr is a LogQL expression filtering logs and returning metric samples.
+//
+//sumtype:decl
 type SampleExpr interface {
 	// Selector is the LogQL selector to apply when retrieving logs.
 	Selector() (LogSelectorExpr, error)
 	Extractor() (SampleExtractor, error)
 	MatcherGroups() ([]MatcherRange, error)
 	Expr
+	isSampleExpr()
 }
 
 // RangeAggregationExpr not all range vector aggregation expressions support grouping by/without label(s),
@@ -1150,6 +1181,7 @@ func newRangeAggregationExpr(left *LogRange, operation string, gr *Grouping, str
 	}
 	return e
 }
+func (e *RangeAggregationExpr) isSampleExpr() {}
 
 func (e *RangeAggregationExpr) Selector() (LogSelectorExpr, error) {
 	if e.err != nil {
@@ -1319,6 +1351,8 @@ func mustNewVectorAggregationExpr(left SampleExpr, operation string, gr *Groupin
 	}
 }
 
+func (e *VectorAggregationExpr) isSampleExpr() {}
+
 func (e *VectorAggregationExpr) MatcherGroups() ([]MatcherRange, error) {
 	if e.err != nil {
 		return nil, e.err
@@ -1880,6 +1914,8 @@ func (e *LiteralExpr) String() string {
 // literlExpr impls SampleExpr & LogSelectorExpr mainly to reduce the need for more complicated typings
 // to facilitate sum types. We'll be type switching when evaluating them anyways
 // and they will only be present in binary operation legs.
+func (e *LiteralExpr) isSampleExpr()                           {}
+func (e *LiteralExpr) isLogSelectorExpr()                      {}
 func (e *LiteralExpr) Selector() (LogSelectorExpr, error)      { return e, e.err }
 func (e *LiteralExpr) HasFilter() bool                         { return false }
 func (e *LiteralExpr) Shardable() bool                         { return true }
@@ -1945,6 +1981,8 @@ func mustNewLabelReplaceExpr(left SampleExpr, dst, replacement, src, regex strin
 	}
 }
 
+func (e *LabelReplaceExpr) isSampleExpr() {}
+
 func (e *LabelReplaceExpr) Selector() (LogSelectorExpr, error) {
 	if e.err != nil {
 		return nil, e.err
@@ -2078,6 +2116,9 @@ func NewVectorExpr(scalar string) *VectorExpr {
 	}
 }
 
+func (e *VectorExpr) isSampleExpr()      {}
+func (e *VectorExpr) isLogSelectorExpr() {}
+
 func (e *VectorExpr) Err() error {
 	return e.err
 }

pkg/logql/syntax/walk.go

@@ -11,3 +11,130 @@ func walkAll(f WalkFn, xs ...Walkable) {
 type Walkable interface {
 	Walk(f WalkFn)
 }
+
+type Visitor interface {
+	SampleExprVisitor
+	LogSelectorExprVisitor
+	StageExprVisitor
+}
+
+type SampleExprVisitor interface {
+	VisitBinOp(*BinOpExpr)
+	VisitVectorAggregation(*VectorAggregationExpr)
+	VisitRangeAggregation(*RangeAggregationExpr)
+	VisitLabelReplace(*LabelReplaceExpr)
+	VisitLiteral(*LiteralExpr)
+	VisitVector(*VectorExpr)
+}
+
+type LogSelectorExprVisitor interface {
+	VisitMatchers(*MatchersExpr)
+	VisitPipeline(*PipelineExpr)
+	VisitLiteral(*LiteralExpr)
+	VisitVector(*VectorExpr)
+}
+
+type StageExprVisitor interface {
+	VisitDecolorize(*DecolorizeExpr)
+	VisitDropLabels(*DropLabelsExpr)
+	VisitJSONExpressionParser(*JSONExpressionParser)
+	VisitKeekLabel(*KeepLabelsExpr)
+	VisitLabelFilter(*LabelFilterExpr)
+	VisitLabelFmt(*LabelFmtExpr)
+	VisitLabelParser(*LabelParserExpr)
+	VisitLineFilter(*LineFilterExpr)
+	VisitLineFmt(*LineFmtExpr)
+	VisitLogfmtExpressionParser(*LogfmtExpressionParser)
+	VisitLogfmtParser(*LogfmtParserExpr)
+}
+
+func Dispatch(expr Expr, v Visitor) {
+	switch e := expr.(type) {
+	case SampleExpr:
+		dispatchSampleExpr(e, v)
+	case LogSelectorExpr:
+		dispatchLogSelectorExpr(e, v)
+	case StageExpr:
+		dispatchStageExpr(e, v)
+		/*
+				case *LogRange
+			case *DecolorizeExpr:
+				v.VisitDecolorize(e)
+			case *DropLabelsExpr:
+				v.VisitDropLabels(e)
+			case *JSONExpressionParser:
+				v.VisitJSONParser(e)
+			case *KeepLabelsExpr:
+				v.VisitKeepLabels(e)
+			case *LabelFilterExpr:
+				v.VisitLabelFilter(e)
+			case *LabelFmtExpr:
+				v.VisitLabelFmt(e)
+			case *LabelParserExpr:
+				v.VisitLabelParser(e)
+			case *LabelReplaceExpr:
+				v.VisitLabelReplace(e)
+			case *LineFilterExpr:
+				v.VisitLineFilter()
+			//LineFmtExpr, LiteralExpr, StageExpr
+		*/
+	}
+}
+
+func dispatchSampleExpr(expr SampleExpr, v SampleExprVisitor) {
+	switch e := expr.(type) {
+	case *BinOpExpr:
+		v.VisitBinOp(e)
+	case *VectorAggregationExpr:
+		v.VisitVectorAggregation(e)
+	case *RangeAggregationExpr:
+		v.VisitRangeAggregation(e)
+	case *LabelReplaceExpr:
+		v.VisitLabelReplace(e)
+	case *LiteralExpr:
+		v.VisitLiteral(e)
+	case *VectorExpr:
+		v.VisitVector(e)
+	}
+}
+
+func dispatchLogSelectorExpr(expr LogSelectorExpr, v LogSelectorExprVisitor) {
+	switch e := expr.(type) {
+	case *PipelineExpr:
+		v.VisitPipeline(e)
+	case *MatchersExpr:
+		v.VisitMatchers(e)
+	case *VectorExpr:
+		v.VisitVector(e)
+	case *LiteralExpr:
+		v.VisitLiteral(e)
+	}
+}
+
+func dispatchStageExpr(expr StageExpr, v StageExprVisitor) {
+	switch e := expr.(type) {
+	case *DecolorizeExpr:
+		v.VisitDecolorize(e)
+	case *DropLabelsExpr:
+		v.VisitDropLabels(e)
+	case *JSONExpressionParser:
+		v.VisitJSONExpressionParser(e)
+	case *KeepLabelsExpr:
+		v.VisitKeekLabel(e)
+	case *LabelFilterExpr:
+		v.VisitLabelFilter(e)
+	case *LabelFmtExpr:
+		v.VisitLabelFmt(e)
+	case *LabelParserExpr:
+		v.VisitLabelParser(e)
+	case *LineFilterExpr:
+		v.VisitLineFilter(e)
+	case *LineFmtExpr:
+		v.VisitLineFmt(e)
+	case *LogfmtExpressionParser:
+		v.VisitLogfmtExpressionParser(e)
+	case *LogfmtParserExpr:
+		v.VisitLogfmtParser(e)
+	}
+
+}