From 20ca79484c36473e14fe7e3d645d98c01262414d Mon Sep 17 00:00:00 2001
From: James Luck <james.luck@thetradedesk.com>
Date: Thu, 25 Nov 2021 07:42:11 +1100
Subject: [PATCH] - Fixing a bug that meant `Duration` was considered a valid
 `Expr`

- Introduced a depth-first visitor that can extract all descendant `Expr` nodes from a given `Expr`
---
 src/PromQL.Parser/Ast.cs                      |  3 +-
 .../DepthFirstExpressionVisitor.cs            | 88 +++++++++++++++++++
 src/PromQL.Parser/IVisitor.cs                 |  2 +-
 .../DepthFirstExpressionVisitorTests.cs       | 51 +++++++++++
 4 files changed, 141 insertions(+), 3 deletions(-)
 create mode 100644 src/PromQL.Parser/DepthFirstExpressionVisitor.cs
 create mode 100644 tests/PromQL.Parser.Tests/DepthFirstExpressionVisitorTests.cs

diff --git a/src/PromQL.Parser/Ast.cs b/src/PromQL.Parser/Ast.cs
index d8c3bf7..047c5f4 100644
--- a/src/PromQL.Parser/Ast.cs
+++ b/src/PromQL.Parser/Ast.cs
@@ -19,7 +19,6 @@ public interface IPromQlNode
     [Closed(
         typeof(AggregateExpr),
         typeof(BinaryExpr),
-        typeof(Duration),
         typeof(FunctionCall),
         typeof(MatrixSelector),
         typeof(NumberLiteral),
@@ -174,7 +173,7 @@ public record NumberLiteral(double Value) : Expr
         public void Accept(IVisitor visitor) => visitor.Visit(this);
     }
 
-    public record Duration(TimeSpan Value) : Expr
+    public record Duration(TimeSpan Value) : IPromQlNode 
     {
         public void Accept(IVisitor visitor) => visitor.Visit(this);
     }
diff --git a/src/PromQL.Parser/DepthFirstExpressionVisitor.cs b/src/PromQL.Parser/DepthFirstExpressionVisitor.cs
new file mode 100644
index 0000000..af96e62
--- /dev/null
+++ b/src/PromQL.Parser/DepthFirstExpressionVisitor.cs
@@ -0,0 +1,88 @@
+using System.Collections;
+using System.Collections.Generic;
+using PromQL.Parser.Ast;
+
+namespace PromQL.Parser
+{
+    /// <summary>
+    /// A depth-first visitor that can find all descendant <see cref="Expr"/> nodes from a given <see cref="Expr"/>.
+    /// </summary>
+    public class DepthFirstExpressionVisitor : IVisitor
+    {
+        private List<Expr> _expressions = new();
+
+        void IVisitor.Visit(StringLiteral expr) => _expressions.Add(expr);
+
+        void IVisitor.Visit(SubqueryExpr sq)
+        {
+            _expressions.Add(sq);
+            sq.Expr.Accept(this);
+        }
+
+        void IVisitor.Visit(Duration d) { }
+
+        void IVisitor.Visit(NumberLiteral n) => _expressions.Add(n);
+
+        void IVisitor.Visit(MetricIdentifier mi) { }
+
+        void IVisitor.Visit(LabelMatcher expr) { }
+
+        void IVisitor.Visit(UnaryExpr unary)
+        {
+            _expressions.Add(unary);
+            unary.Expr.Accept(this);
+        }
+
+        void IVisitor.Visit(MatrixSelector ms)
+        {
+            _expressions.Add(ms);
+            // No need to visit vector selector, it's accessible from matrix selector
+        }
+
+        void IVisitor.Visit(OffsetExpr offset)
+        {
+            _expressions.Add(offset);
+            offset.Expr.Accept(this);
+        }
+
+        void IVisitor.Visit(ParenExpression paren)
+        {
+            _expressions.Add(paren);
+            paren.Expr.Accept(this);
+        }
+
+        void IVisitor.Visit(FunctionCall fnCall)
+        {
+            _expressions.Add(fnCall);
+            foreach (var a in fnCall.Args)
+                a.Accept(this);
+        }
+
+        void IVisitor.Visit(VectorMatching vm) { }
+
+        void IVisitor.Visit(BinaryExpr expr)
+        {
+            _expressions.Add(expr);
+            expr.LeftHandSide.Accept(this);
+            expr.RightHandSide.Accept(this);
+        }
+
+        void IVisitor.Visit(AggregateExpr expr)
+        {
+            _expressions.Add(expr);
+            expr.Param?.Accept(this);
+            expr.Expr.Accept(this);
+        }
+
+        void IVisitor.Visit(VectorSelector vs) => _expressions.Add(vs);
+
+        void IVisitor.Visit(LabelMatchers lms) { }
+
+        public IEnumerable<Expr> GetExpressions(Expr expr)
+        {
+            _expressions.Clear();
+            expr.Accept(this);
+            return _expressions;
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/PromQL.Parser/IVisitor.cs b/src/PromQL.Parser/IVisitor.cs
index 35962e3..e71cc3d 100644
--- a/src/PromQL.Parser/IVisitor.cs
+++ b/src/PromQL.Parser/IVisitor.cs
@@ -27,6 +27,6 @@ public interface IVisitor
         void Visit(BinaryExpr expr);
         void Visit(AggregateExpr expr);
         void Visit(VectorSelector vs);
-        void Visit(LabelMatchers fnCall);
+        void Visit(LabelMatchers lms);
     }
 }
\ No newline at end of file
diff --git a/tests/PromQL.Parser.Tests/DepthFirstExpressionVisitorTests.cs b/tests/PromQL.Parser.Tests/DepthFirstExpressionVisitorTests.cs
new file mode 100644
index 0000000..3ee4c42
--- /dev/null
+++ b/tests/PromQL.Parser.Tests/DepthFirstExpressionVisitorTests.cs
@@ -0,0 +1,51 @@
+using System.Linq;
+using FluentAssertions;
+using NUnit.Framework;
+using PromQL.Parser.Ast;
+
+namespace PromQL.Parser.Tests
+{
+    [TestFixture]
+    public class DepthFirstExpressionVisitorTests
+    {
+        [Test]
+        public void Visit_Basic_Types()
+        {
+            // Not semantically valid but syntactically valid!
+            const string toParse = "'hello' + 1";
+            var expr = Parser.ParseExpression(toParse);
+            var visitor = new DepthFirstExpressionVisitor();
+            visitor.GetExpressions(expr)
+                .Select(x => x.GetType())
+                .Should()
+                .Equal(
+                    typeof(BinaryExpr),
+                    typeof(StringLiteral),
+                    typeof(NumberLiteral)
+                );
+        }
+        
+        [Test]
+        public void Visit_Complex_Expression()
+        {
+            const string toParse = "sum(rate(my_vector[1m] offset 5m)) + -(some_metric[1m:])";
+            var expr = Parser.ParseExpression(toParse);
+            var visitor = new DepthFirstExpressionVisitor();
+            visitor.GetExpressions(expr)
+                    .Select(x => x.GetType())
+                    .Should()
+                    .Equal(
+                        typeof(BinaryExpr),
+                        typeof(AggregateExpr),
+                        typeof(FunctionCall),
+                        typeof(OffsetExpr),
+                        typeof(MatrixSelector),
+                        typeof(UnaryExpr),
+                        typeof(ParenExpression),
+                        typeof(SubqueryExpr),
+                        typeof(VectorSelector)
+                    )
+                ;
+        }
+    }
+}
\ No newline at end of file