feat(evaluator): implement graph builtins

opa-builtins/README.md

@@ -94,6 +94,8 @@ Note: String builtins (`contains`, `concat`, `split`, `sprintf`, `trim`, etc.) a
 | `net.cidr_contains`, `net.cidr_contains_matches` | Yes |
 | `net.cidr_intersects`, `net.cidr_expand`, `net.cidr_merge` | Yes |
 | `net.cidr_is_valid`, `net.lookup_ip_addr` | Yes |
+| **Graphs** (opa-evaluator) | |
+| `graph.reachable`, `graph.reachable_paths` | Yes |
 | **Semantic Versions** (opa-builtins-semver) | |
 | `semver.compare`, `semver.is_valid` | Yes |
 | **Providers** (opa-builtins-providers-aws) | |
@@ -103,7 +105,7 @@ Note: String builtins (`contains`, `concat`, `split`, `sprintf`, `trim`, etc.) a
 | **Comparison** | |
 | `equal`, `neq`, `lt`, `lte`, `gt`, `gte` | Yes |
 | **Not Yet Implemented** | |
-| `bits.*`, `graph.*`, `units.*`, `http.send` | No |
+| `bits.*`, `units.*`, `http.send` | No |
 | `uuid.*`, `graphql.*`, `rego.*` | No |
 
 ## Adding Custom Builtins

opa-evaluator/src/main/java/io/github/open_policy_agent/opa/ast/builtin/BuiltinRegistry.java

@@ -14,6 +14,7 @@
 import io.github.open_policy_agent.opa.ast.builtin.impls.CastBuiltins;
 import io.github.open_policy_agent.opa.ast.builtin.impls.ComparisonBuiltins;
 import io.github.open_policy_agent.opa.ast.builtin.impls.EncodingBuiltins;
+import io.github.open_policy_agent.opa.ast.builtin.impls.GraphBuiltins;
 import io.github.open_policy_agent.opa.ast.builtin.impls.HexBuiltins;
 import io.github.open_policy_agent.opa.ast.builtin.impls.ObjectBuiltins;
 import io.github.open_policy_agent.opa.ast.builtin.impls.OpaBuiltins;
@@ -40,6 +41,7 @@ public class BuiltinRegistry {
     ArithmeticBuiltins.class,
     ArrayBuiltins.class,
     EncodingBuiltins.class,
+    GraphBuiltins.class,
     HexBuiltins.class,
     UriBuiltins.class,
     StringBuiltins.class,

opa-evaluator/src/main/java/io/github/open_policy_agent/opa/ast/builtin/impls/GraphBuiltins.java

@@ -0,0 +1,192 @@
+package io.github.open_policy_agent.opa.ast.builtin.impls;
+
+import static io.github.open_policy_agent.opa.ast.builtin.impls.utils.ArgHelper.getArg;
+
+import io.github.open_policy_agent.opa.ast.builtin.OpaBuiltin;
+import io.github.open_policy_agent.opa.ast.builtin.OpaDynamic;
+import io.github.open_policy_agent.opa.ast.builtin.OpaType;
+import io.github.open_policy_agent.opa.ast.builtin.OpaVal;
+import io.github.open_policy_agent.opa.ast.types.RegoArray;
+import io.github.open_policy_agent.opa.ast.types.RegoCollection;
+import io.github.open_policy_agent.opa.ast.types.RegoObject;
+import io.github.open_policy_agent.opa.ast.types.RegoSet;
+import io.github.open_policy_agent.opa.ast.types.RegoValue;
+import io.github.open_policy_agent.opa.rego.EvaluationContext;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.HashSet;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
+
+public class GraphBuiltins {
+
+  public static Map<String, BiFunction<EvaluationContext, RegoValue[], RegoValue>> builtins() {
+    GraphBuiltins instance = new GraphBuiltins();
+    return Map.of(
+        "graph.reachable", instance::reachable,
+        "graph.reachable_paths", instance::reachablePaths);
+  }
+
+  @OpaBuiltin(
+      name = "graph.reachable",
+      description =
+          "Computes the set of reachable nodes in the graph from a set of starting nodes.",
+      categories = {"graphs"},
+      args = {
+        @OpaType(
+            type = "object",
+            name = "graph",
+            description = "object containing a set or array of neighboring vertices",
+            dynamic = @OpaDynamic(keyType = "any", valueType = "any")),
+        @OpaType(
+            name = "initial",
+            description = "set or array of root vertices",
+            of = {@OpaVal("set"), @OpaVal("array")})
+      },
+      result =
+          @OpaType(
+              type = "set",
+              name = "output",
+              description = "vertices reachable from the initial vertices in the directed graph",
+              dynamic = @OpaDynamic(type = "any")))
+  public RegoSet reachable(EvaluationContext ctx, RegoValue[] args) {
+    RegoObject graph = getArg(args, 0, RegoObject.class);
+    RegoCollection initial = getArg(args, 1, RegoCollection.class);
+    Deque<RegoValue> queue = new ArrayDeque<>(initial.valueStream().toList());
+    RegoSet reached = new RegoSet(false);
+
+    while (!queue.isEmpty()) {
+      RegoValue node = queue.removeFirst();
+      RegoValue edges = graph.getProperty(node);
+      if (edges == null) {
+        continue;
+      }
+      for (RegoValue neighbor : collectionValues(edges)) {
+        if (!reached.contains(neighbor)) {
+          queue.addLast(neighbor);
+        }
+      }
+      reached.addValue(node);
+    }
+
+    return new RegoSet(ctx.sortSets, reached.getValue());
+  }
+
+  @OpaBuiltin(
+      name = "graph.reachable_paths",
+      description =
+          "Computes the set of reachable paths in the graph from a set of starting nodes.",
+      categories = {"graphs"},
+      args = {
+        @OpaType(
+            type = "object",
+            name = "graph",
+            description = "object containing a set or array of neighboring vertices",
+            dynamic = @OpaDynamic(keyType = "any", valueType = "any")),
+        @OpaType(
+            name = "initial",
+            description = "set or array of root vertices",
+            of = {@OpaVal("set"), @OpaVal("array")})
+      },
+      result =
+          @OpaType(
+              type = "set",
+              name = "output",
+              description = "paths reachable from the initial vertices in the directed graph",
+              dynamic = @OpaDynamic(type = "array")))
+  public RegoSet reachablePaths(EvaluationContext ctx, RegoValue[] args) {
+    RegoObject graph = getArg(args, 0, RegoObject.class);
+    RegoCollection initial = getArg(args, 1, RegoCollection.class);
+    RegoSet paths = new RegoSet(false);
+
+    for (RegoValue node : initial.valueStream().toList()) {
+      collectPaths(graph, node, paths);
+    }
+
+    if (!ctx.sortSets) {
+      return paths;
+    }
+    List<RegoValue> sortedPaths = new ArrayList<>(paths.getValue());
+    sortedPaths.sort(RegoValue::compareTo);
+    return new RegoSet(false, new LinkedHashSet<>(sortedPaths));
+  }
+
+  private static void collectPaths(RegoObject graph, RegoValue initial, RegoSet paths) {
+    RegoValue edges = graph.getProperty(initial);
+    if (edges == null) {
+      return;
+    }
+
+    List<RegoValue> neighbors = collectionValues(edges);
+    if (neighbors.isEmpty()) {
+      addPath(paths, List.of(initial));
+      return;
+    }
+
+    List<RegoValue> path = new ArrayList<>(List.of(initial));
+    Set<RegoValue> reached = new HashSet<>(Set.of(initial));
+    Deque<PathFrame> pending = new ArrayDeque<>();
+    pending.addLast(new PathFrame(null, neighbors));
+
+    while (!pending.isEmpty()) {
+      PathFrame frame = pending.peekLast();
+      if (frame.nextNeighborIndex == frame.neighbors.size()) {
+        pending.removeLast();
+        if (frame.node != null) {
+          path.remove(path.size() - 1);
+          reached.remove(frame.node);
+        }
+        continue;
+      }
+
+      RegoValue neighbor = frame.neighbors.get(frame.nextNeighborIndex++);
+      if (reached.contains(neighbor)) {
+        addPath(paths, path);
+        continue;
+      }
+
+      RegoValue neighborEdges = graph.getProperty(neighbor);
+      if (neighborEdges == null) {
+        addPath(paths, path);
+        continue;
+      }
+
+      path.add(neighbor);
+      reached.add(neighbor);
+      List<RegoValue> neighborValues = collectionValues(neighborEdges);
+      if (neighborValues.isEmpty()) {
+        addPath(paths, path);
+        path.remove(path.size() - 1);
+        reached.remove(neighbor);
+        continue;
+      }
+      pending.addLast(new PathFrame(neighbor, neighborValues));
+    }
+  }
+
+  private static void addPath(RegoSet paths, List<RegoValue> path) {
+    paths.addValue(new RegoArray(new ArrayList<>(path)));
+  }
+
+  private static final class PathFrame {
+    private final RegoValue node;
+    private final List<RegoValue> neighbors;
+    private int nextNeighborIndex;
+
+    private PathFrame(RegoValue node, List<RegoValue> neighbors) {
+      this.node = node;
+      this.neighbors = neighbors;
+    }
+  }
+
+  private static List<RegoValue> collectionValues(RegoValue value) {
+    if (value instanceof RegoCollection) {
+      return ((RegoCollection) value).valueStream().toList();
+    }
+    return List.of();
+  }
+}

opa-evaluator/src/test/java/io/github/open_policy_agent/opa/ast/builtin/impls/GraphBuiltinsTest.java

@@ -0,0 +1,93 @@
+package io.github.open_policy_agent.opa.ast.builtin.impls;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import io.github.open_policy_agent.opa.ast.builtin.BuiltinRegistry;
+import io.github.open_policy_agent.opa.ast.types.RegoArray;
+import io.github.open_policy_agent.opa.ast.types.RegoObject;
+import io.github.open_policy_agent.opa.ast.types.RegoSet;
+import io.github.open_policy_agent.opa.ast.types.RegoString;
+import io.github.open_policy_agent.opa.ast.types.RegoValue;
+import io.github.open_policy_agent.opa.rego.EvaluationContext;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiFunction;
+import org.junit.jupiter.api.Test;
+
+class GraphBuiltinsTest {
+
+  private static final EvaluationContext CONTEXT = new EvaluationContext.Builder().build();
+
+  @Test
+  void reachableTraversesArrayAndSetEdges() {
+    RegoString a = new RegoString("a");
+    RegoString b = new RegoString("b");
+    RegoString c = new RegoString("c");
+    RegoString d = new RegoString("d");
+    RegoObject graph =
+        new RegoObject(
+            Map.of(
+                a, new RegoArray(List.of(b, c)),
+                b, setOf(d),
+                c, new RegoArray(List.of(d)),
+                d, setOf()));
+
+    RegoValue result = call("graph.reachable", graph, setOf(a, new RegoString("missing")));
+
+    assertEquals(setOf(a, b, c, d), result);
+  }
+
+  @Test
+  void reachablePathsStopsAtCyclesAndMissingNodes() {
+    RegoString a = new RegoString("a");
+    RegoString b = new RegoString("b");
+    RegoString missing = new RegoString("missing");
+    RegoObject graph =
+        new RegoObject(
+            Map.of(
+                a, new RegoArray(List.of(b, missing)),
+                b, setOf(a)));
+
+    RegoValue result = call("graph.reachable_paths", graph, setOf(a));
+
+    assertEquals(
+        setOf(new RegoArray(List.of(a, b)), new RegoArray(List.of(a))), result);
+  }
+
+  @Test
+  void reachablePathsHandlesDeepGraphsWithoutRecursion() {
+    int nodeCount = 10_000;
+    List<RegoValue> nodes = new ArrayList<>(nodeCount);
+    Map<RegoValue, RegoValue> graph = new LinkedHashMap<>();
+
+    for (int index = 0; index < nodeCount; index++) {
+      nodes.add(new RegoString("node-" + index));
+    }
+    for (int index = 0; index < nodeCount - 1; index++) {
+      graph.put(nodes.get(index), new RegoArray(List.of(nodes.get(index + 1))));
+    }
+    graph.put(nodes.get(nodeCount - 1), setOf());
+
+    RegoValue result = call("graph.reachable_paths", new RegoObject(graph), setOf(nodes.get(0)));
+
+    assertEquals(setOf(new RegoArray(nodes)), result);
+  }
+
+  private static RegoValue call(String name, RegoValue... args) {
+    BiFunction<EvaluationContext, RegoValue[], RegoValue> builtin =
+        BuiltinRegistry.AllBuiltIns.get(name);
+    assertNotNull(builtin, name + " should be registered");
+    return builtin.apply(CONTEXT, args);
+  }
+
+  private static RegoSet setOf(RegoValue... values) {
+    RegoSet set = new RegoSet(false);
+    for (RegoValue value : values) {
+      set.addValue(value);
+    }
+    return set;
+  }
+}