Add the ability to navigate ASTs

PiperOrigin-RevId: 558256662

common/navigation/BUILD.bazel

@@ -0,0 +1,9 @@
+package(
+    default_applicable_licenses = ["//:license"],
+    default_visibility = ["//visibility:public"],  # TODO: Expose when ready
+)
+
+java_library(
+    name = "navigation",
+    exports = ["//common/src/main/java/dev/cel/common/navigation"],
+)

common/src/main/java/dev/cel/common/navigation/BUILD.bazel

@@ -0,0 +1,25 @@
+package(
+    default_applicable_licenses = [
+        "//:license",
+    ],
+    default_visibility = [
+        "//common/navigation:__pkg__",
+    ],
+)
+
+java_library(
+    name = "navigation",
+    srcs = [
+        "CelNavigableAst.java",
+        "CelNavigableExpr.java",
+        "CelNavigableExprVisitor.java",
+    ],
+    tags = [
+    ],
+    deps = [
+        "//:auto_value",
+        "//common",
+        "//common/ast",
+        "@maven//:com_google_guava_guava",
+    ],
+)

common/src/main/java/dev/cel/common/navigation/CelNavigableAst.java

@@ -0,0 +1,47 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.navigation;
+
+import dev.cel.common.CelAbstractSyntaxTree;
+
+/**
+ * Decorates a {@link CelAbstractSyntaxTree} with navigational properties. This allows us to visit a
+ * node's children, descendants or its parent with ease.
+ */
+public final class CelNavigableAst {
+
+  private final CelAbstractSyntaxTree ast;
+  private final CelNavigableExpr root;
+
+  private CelNavigableAst(CelAbstractSyntaxTree ast) {
+    this.ast = ast;
+    this.root = CelNavigableExpr.builder().setExpr(ast.getExpr()).setDepth(0).build();
+  }
+
+  /** Constructs a new instance of {@link CelNavigableAst} from {@link CelAbstractSyntaxTree}. */
+  public static CelNavigableAst fromAst(CelAbstractSyntaxTree ast) {
+    return new CelNavigableAst(ast);
+  }
+
+  /** Returns the root of the AST. */
+  public CelNavigableExpr getRoot() {
+    return root;
+  }
+
+  /** Returns the underlying {@link CelAbstractSyntaxTree}. */
+  public CelAbstractSyntaxTree getAst() {
+    return ast;
+  }
+}

common/src/main/java/dev/cel/common/navigation/CelNavigableExpr.java

@@ -0,0 +1,114 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.navigation;
+
+import com.google.auto.value.AutoValue;
+import dev.cel.common.ast.CelExpr;
+import dev.cel.common.ast.CelExpr.CelComprehension;
+import dev.cel.common.ast.CelExpr.ExprKind;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+/**
+ * CelNavigableExpr represents the base navigable expression value with methods to inspect the
+ * parent and child expressions.
+ */
+@AutoValue
+public abstract class CelNavigableExpr {
+
+  /**
+   * Specifies the traversal order of AST navigation.
+   *
+   * <p>For call expressions, the target is visited before its arguments.
+   *
+   * <p>For comprehensions, the visiting order is as follows:
+   *
+   * <ol>
+   *   <li>{@link CelComprehension#iterRange}
+   *   <li>{@link CelComprehension#accuInit}
+   *   <li>{@link CelComprehension#loopCondition}
+   *   <li>{@link CelComprehension#loopStep}
+   *   <li>{@link CelComprehension#result}
+   * </ol>
+   */
+  public enum TraversalOrder {
+    PRE_ORDER,
+    POST_ORDER
+  }
+
+  public abstract CelExpr expr();
+
+  public abstract Optional<CelNavigableExpr> parent();
+
+  /** Represents the count of transitive parents. Depth of an AST's root is 0. */
+  public abstract int depth();
+
+  /**
+   * Returns a stream of {@link CelNavigableExpr} collected from the current node down to the last
+   * leaf-level member using post-order traversal.
+   */
+  public Stream<CelNavigableExpr> descendants() {
+    return descendants(TraversalOrder.POST_ORDER);
+  }
+
+  /**
+   * Returns a stream of {@link CelNavigableExpr} collected from the current node down to the last
+   * leaf-level member using the specified traversal order.
+   */
+  public Stream<CelNavigableExpr> descendants(TraversalOrder traversalOrder) {
+    return CelNavigableExprVisitor.collect(this, traversalOrder);
+  }
+
+  /**
+   * Returns a stream of {@link CelNavigableExpr} collected from the current node to its immediate
+   * children using post-order traversal.
+   */
+  public Stream<CelNavigableExpr> children() {
+    return children(TraversalOrder.POST_ORDER);
+  }
+
+  /**
+   * Returns a stream of {@link CelNavigableExpr} collected from the current node to its immediate
+   * children using the specified traversal order.
+   */
+  public Stream<CelNavigableExpr> children(TraversalOrder traversalOrder) {
+    return CelNavigableExprVisitor.collect(this, 1, traversalOrder);
+  }
+
+  /** Returns the underlying kind of the {@link CelExpr}. */
+  public ExprKind.Kind getKind() {
+    return expr().exprKind().getKind();
+  }
+
+  /** Create a new builder to construct a {@link CelNavigableExpr} instance. */
+  public static Builder builder() {
+    return new AutoValue_CelNavigableExpr.Builder();
+  }
+
+  /** Builder to configure {@link CelNavigableExpr}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    public abstract CelExpr expr();
+
+    public abstract Builder setExpr(CelExpr value);
+
+    public abstract Builder setParent(CelNavigableExpr value);
+
+    public abstract Builder setDepth(int value);
+
+    public abstract CelNavigableExpr build();
+  }
+}

common/src/main/java/dev/cel/common/navigation/CelNavigableExprVisitor.java

@@ -0,0 +1,196 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.navigation;
+
+import com.google.common.collect.ImmutableList;
+import dev.cel.common.ast.CelExpr;
+import dev.cel.common.ast.CelExpr.CelCall;
+import dev.cel.common.ast.CelExpr.CelComprehension;
+import dev.cel.common.ast.CelExpr.CelCreateList;
+import dev.cel.common.ast.CelExpr.CelCreateStruct;
+import dev.cel.common.ast.CelExpr.CelCreateStruct.Entry;
+import dev.cel.common.ast.CelExpr.CelCreateStruct.Entry.KeyKind.Kind;
+import dev.cel.common.ast.CelExpr.CelSelect;
+import dev.cel.common.navigation.CelNavigableExpr.TraversalOrder;
+import java.util.stream.Stream;
+
+/** Visitor implementation to navigate an AST. */
+final class CelNavigableExprVisitor {
+
+  private static final int MAX_DESCENDANTS_RECURSION_DEPTH = 500;
+
+  private final Stream.Builder<CelNavigableExpr> streamBuilder;
+  private final TraversalOrder traversalOrder;
+  private final int maxDepth;
+
+  private CelNavigableExprVisitor(int maxDepth, TraversalOrder traversalOrder) {
+    this.maxDepth = maxDepth;
+    this.traversalOrder = traversalOrder;
+    this.streamBuilder = Stream.builder();
+  }
+
+  /**
+   * Returns a stream containing all the nodes using the specified traversal order. Note that the
+   * collection occurs eagerly.
+   *
+   * <pre>
+   * For example, given the following tree:
+   *             a
+   *          b     c
+   *       d    e
+   *
+   * The collected nodes are as follows using pre-order traversal:
+   *
+   * maxDepth of -1: none
+   * maxDepth of 0: a
+   * maxDepth of 1: a, b, c
+   * maxDepth of 2: a, b, d, e, c
+   * </pre>
+   */
+  static Stream<CelNavigableExpr> collect(
+      CelNavigableExpr navigableExpr, TraversalOrder traversalOrder) {
+    return collect(navigableExpr, MAX_DESCENDANTS_RECURSION_DEPTH, traversalOrder);
+  }
+
+  /**
+   * Returns a stream containing all the nodes upto the maximum depth specified using the specified
+   * traversal order. Note that the collection occurs eagerly.
+   *
+   * <pre>
+   * For example, given the following tree:
+   *             a
+   *          b     c
+   *       d    e
+   *
+   * The collected nodes are as follows using pre-order traversal:
+   *
+   * maxDepth of -1: none
+   * maxDepth of 0: a
+   * maxDepth of 1: a, b, c
+   * maxDepth of 2: a, b, d, e, c
+   * </pre>
+   */
+  static Stream<CelNavigableExpr> collect(
+      CelNavigableExpr navigableExpr, int maxDepth, TraversalOrder traversalOrder) {
+    CelNavigableExprVisitor visitor = new CelNavigableExprVisitor(maxDepth, traversalOrder);
+
+    visitor.visit(navigableExpr);
+
+    return visitor.streamBuilder.build();
+  }
+
+  private void visit(CelNavigableExpr navigableExpr) {
+    if (navigableExpr.depth() > MAX_DESCENDANTS_RECURSION_DEPTH - 1) {
+      throw new IllegalStateException("Max recursion depth reached.");
+    }
+    if (navigableExpr.depth() > maxDepth) {
+      return;
+    }
+    if (traversalOrder.equals(TraversalOrder.PRE_ORDER)) {
+      streamBuilder.add(navigableExpr);
+    }
+
+    switch (navigableExpr.getKind()) {
+      case CALL:
+        visit(navigableExpr, navigableExpr.expr().call());
+        break;
+      case CREATE_LIST:
+        visit(navigableExpr, navigableExpr.expr().createList());
+        break;
+      case SELECT:
+        visit(navigableExpr, navigableExpr.expr().select());
+        break;
+      case CREATE_STRUCT:
+        CelCreateStruct createStruct = navigableExpr.expr().createStruct();
+        if (createStruct.entries().isEmpty()) {
+          break;
+        }
+        // TODO: Replace with CREATE_MAP kind
+        if (createStruct.entries().get(0).keyKind().getKind().equals(Kind.MAP_KEY)) {
+          visitMap(navigableExpr, navigableExpr.expr().createStruct());
+        } else {
+          visitStruct(navigableExpr, navigableExpr.expr().createStruct());
+        }
+        break;
+      case COMPREHENSION:
+        visit(navigableExpr, navigableExpr.expr().comprehension());
+        break;
+      default:
+        break;
+    }
+
+    if (traversalOrder.equals(TraversalOrder.POST_ORDER)) {
+      streamBuilder.add(navigableExpr);
+    }
+  }
+
+  private void visit(CelNavigableExpr navigableExpr, CelCall call) {
+    if (call.target().isPresent()) {
+      CelNavigableExpr target = newNavigableChild(navigableExpr, call.target().get());
+      visit(target);
+    }
+
+    visitExprList(call.args(), navigableExpr);
+  }
+
+  private void visit(CelNavigableExpr navigableExpr, CelCreateList createList) {
+    visitExprList(createList.elements(), navigableExpr);
+  }
+
+  private void visit(CelNavigableExpr navigableExpr, CelSelect selectExpr) {
+    CelNavigableExpr operand = newNavigableChild(navigableExpr, selectExpr.operand());
+    visit(operand);
+  }
+
+  private void visit(CelNavigableExpr navigableExpr, CelComprehension comprehension) {
+    visit(newNavigableChild(navigableExpr, comprehension.iterRange()));
+    visit(newNavigableChild(navigableExpr, comprehension.accuInit()));
+    visit(newNavigableChild(navigableExpr, comprehension.loopCondition()));
+    visit(newNavigableChild(navigableExpr, comprehension.loopStep()));
+    visit(newNavigableChild(navigableExpr, comprehension.result()));
+  }
+
+  private void visitStruct(CelNavigableExpr navigableExpr, CelCreateStruct struct) {
+    for (Entry entry : struct.entries()) {
+      CelNavigableExpr value = newNavigableChild(navigableExpr, entry.value());
+      visit(value);
+    }
+  }
+
+  private void visitMap(CelNavigableExpr navigableExpr, CelCreateStruct struct) {
+    for (Entry entry : struct.entries()) {
+      CelNavigableExpr key = newNavigableChild(navigableExpr, entry.keyKind().mapKey());
+      visit(key);
+
+      CelNavigableExpr value = newNavigableChild(navigableExpr, entry.value());
+      visit(value);
+    }
+  }
+
+  private void visitExprList(ImmutableList<CelExpr> createListExpr, CelNavigableExpr parent) {
+    for (CelExpr expr : createListExpr) {
+      CelNavigableExpr arg = newNavigableChild(parent, expr);
+      visit(arg);
+    }
+  }
+
+  private CelNavigableExpr newNavigableChild(CelNavigableExpr parent, CelExpr expr) {
+    return CelNavigableExpr.builder()
+        .setExpr(expr)
+        .setDepth(parent.depth() + 1)
+        .setParent(parent)
+        .build();
+  }
+}