From c05a0ddcb7ba25a8802fc611bd4b25f51b439ea1 Mon Sep 17 00:00:00 2001
From: TristonianJones <tswadell@google.com>
Date: Fri, 18 Aug 2023 16:42:13 -0700
Subject: [PATCH] Updates based on feedback

---
 common/ast/navigable.go      |  3 ++
 common/ast/navigable_test.go | 71 ++++++++++++++++++++++++++++--------
 2 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/common/ast/navigable.go b/common/ast/navigable.go
index 6df5d8967..2836b5655 100644
--- a/common/ast/navigable.go
+++ b/common/ast/navigable.go
@@ -188,6 +188,9 @@ const (
 	postOrder
 )
 
+// TODO: consider exposing a way to configure a limit for the max visit depth.
+// It's possible that we could want to configure this on the NewExprVisitor()
+// and through MatchDescendents() / MaxID().
 func visit(expr Expr, visitor Visitor, order visitOrder, depth, maxDepth int) {
 	if maxDepth > 0 && depth == maxDepth {
 		return
diff --git a/common/ast/navigable_test.go b/common/ast/navigable_test.go
index f4b8fa6b8..c4f312df8 100644
--- a/common/ast/navigable_test.go
+++ b/common/ast/navigable_test.go
@@ -137,14 +137,51 @@ func TestNavigateAST(t *testing.T) {
 
 func TestNavigableExprChildren(t *testing.T) {
 	tests := []struct {
-		expr string
+		expr         string
+		preOrderIDs  []int64
+		postOrderIDs []int64
 	}{
-		{expr: `'a' == 'b'`},
-		{expr: `'a'.size()`},
-		{expr: `type(1) == int`},
-		{expr: `{'hello': 'world'}.hello`},
-		{expr: `google.expr.proto3.test.TestAllTypes{single_int32: 1}`},
-		{expr: `[true].exists(i, i)`},
+		{
+			// [2] ==, [1] 'a', [3] 'b'
+			expr:         `'a' == 'b'`,
+			preOrderIDs:  []int64{2, 1, 3},
+			postOrderIDs: []int64{1, 3, 2},
+		},
+		{
+			// [2] size(), [1] 'a'
+			expr:         `'a'.size()`,
+			preOrderIDs:  []int64{2, 1},
+			postOrderIDs: []int64{1, 2},
+		},
+		{
+			// [3] ==, [1] type(), [2] 1, [4] int
+			expr:         `type(1) == int`,
+			preOrderIDs:  []int64{3, 1, 2, 4},
+			postOrderIDs: []int64{2, 1, 4, 3},
+		},
+		{
+			// [5] .hello, [1] {}, [3] 'hello', [4] 'world'
+			expr:         `{'hello': 'world'}.hello`,
+			preOrderIDs:  []int64{5, 1, 3, 4},
+			postOrderIDs: []int64{3, 4, 1, 5},
+		},
+		{
+			// [1] TestAllTypes, [3] 1
+			expr:         `google.expr.proto3.test.TestAllTypes{single_int32: 1}`,
+			preOrderIDs:  []int64{1, 3},
+			postOrderIDs: []int64{3, 1},
+		},
+		{
+			// [13] comprehension
+			// range:    [1] [], [2] true
+			// accuInit: [6] result=false
+			// loopCond: [9] @not_strictly_false, [8] !, [7] result
+			// loopStep: [11] ||, [10] result [5] i
+			// result:   [12] result
+			expr:         `[true].exists(i, i)`,
+			preOrderIDs:  []int64{13, 1, 2, 6, 9, 8, 7, 11, 10, 5, 12},
+			postOrderIDs: []int64{2, 1, 6, 7, 8, 9, 10, 5, 11, 12, 13},
+		},
 	}
 
 	for _, tst := range tests {
@@ -158,16 +195,18 @@ func TestNavigableExprChildren(t *testing.T) {
 				preOrderExprIDs = append(preOrderExprIDs, nav.ID())
 			})
 			ast.PreOrderVisit(root, navVisitor)
-
-			allExprs := []int64{}
-			visited := []ast.NavigableExpr{root}
-			for len(visited) > 0 {
-				e := visited[0]
-				allExprs = append(allExprs, e.ID())
-				visited = append(e.Children()[:], visited[1:]...)
+			if !reflect.DeepEqual(tc.preOrderIDs, preOrderExprIDs) {
+				t.Errorf("PreOrderVisit() got %v expressions, wanted %v", tc.preOrderIDs, preOrderExprIDs)
 			}
-			if !reflect.DeepEqual(allExprs, preOrderExprIDs) {
-				t.Fatalf("Children() got %v expressions, wanted %v", allExprs, preOrderExprIDs)
+
+			postOrderExprIDs := []int64{}
+			navVisitor = ast.NewExprVisitor(func(e ast.Expr) {
+				nav := e.(ast.NavigableExpr)
+				postOrderExprIDs = append(postOrderExprIDs, nav.ID())
+			})
+			ast.PostOrderVisit(root, navVisitor)
+			if !reflect.DeepEqual(tc.postOrderIDs, postOrderExprIDs) {
+				t.Errorf("PostOrderVisit() got %v expressions, wanted %v", tc.postOrderIDs, postOrderExprIDs)
 			}
 		})
 	}