feat(evaluator): implement walk builtin

CHANGELOG.md

@@ -5,6 +5,12 @@ project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+### Runtime, SDK, Tooling
+
+- Implement the `walk` builtin
+- Encode composite object keys (arrays/sets/objects) as compact JSON to match
+  Go-OPA, instead of leaking Java's collection formatting
+
 ### Build and CI
 
 - Fail the compliance suite on fixtures whose builtin cannot be resolved, instead

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;
@@ -39,6 +40,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,94 @@
+package io.github.open_policy_agent.opa.ast.builtin.impls;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiFunction;
+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.types.RegoArray;
+import io.github.open_policy_agent.opa.ast.types.RegoInt32;
+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;
+
+public class GraphBuiltins {
+
+  public static Map<String, BiFunction<EvaluationContext, RegoValue[], RegoValue>> builtins() {
+    GraphBuiltins instance = new GraphBuiltins();
+    return Map.of("walk", instance::walk);
+  }
+
+  /**
+   * {@code walk} is the one builtin Go-OPA flags as relational ({@code Relation: true}): it yields
+   * many {@code [path, value]} tuples through an iterator instead of returning a single value. The
+   * IR planner turns it into an ordinary call whose result is scanned -- the plan emits a CallStmt
+   * for {@code walk} followed by a ScanStmt over the returned collection -- so what is returned
+   * here is the complete set of pairs rather than one tuple.
+   */
+  @OpaBuiltin(
+      name = "walk",
+      description =
+          "Generates `[path, value]` tuples for all nested documents of `x` (recursively).  Queries"
+              + " can use `walk` to traverse documents nested under `x`.",
+      categories = {"graph"},
+      args = {@OpaType(type = "any", name = "x", description = "value to walk")},
+      result =
+          @OpaType(
+              type = "array",
+              name = "output",
+              dynamic = @OpaDynamic(type = "any"),
+              description =
+                  "pairs of `path` and `value`: `path` is an array representing the pointer to"
+                      + " `value` in `x`. If `path` is assigned a wildcard (`_`), the `walk`"
+                      + " function will skip path creation entirely for faster evaluation."))
+  public RegoSet walk(EvaluationContext ctx, RegoValue[] args) {
+    RegoSet pairs = new RegoSet(ctx.sortSets);
+    // An undefined argument produces no tuples, matching the undefined expression in Go-OPA.
+    if (args == null || args.length == 0 || args[0] == null) {
+      return pairs;
+    }
+    collect(new ArrayDeque<>(), args[0], pairs);
+    return pairs;
+  }
+
+  /**
+   * Emits the {@code [path, value]} pair for {@code value} and then recurses into its members. The
+   * root is always emitted with an empty path, including for scalars, which have no members.
+   */
+  private void collect(Deque<RegoValue> path, RegoValue value, RegoSet pairs) {
+    RegoArray pair = new RegoArray(2);
+    // Snapshot the path: the deque is mutated as the walk unwinds, so sharing it would leave every
+    // pair pointing at the same (final) path.
+    pair.addValue(new RegoArray(new ArrayList<>(path)));
+    pair.addValue(value);
+    pairs.addValue(pair);
+
+    if (value instanceof RegoObject) {
+      for (Map.Entry<RegoValue, RegoValue> entry :
+          ((RegoObject) value).getProperties().entrySet()) {
+        path.addLast(entry.getKey());
+        collect(path, entry.getValue(), pairs);
+        path.removeLast();
+      }
+    } else if (value instanceof RegoArray) {
+      List<RegoValue> values = ((RegoArray) value).getValue();
+      for (int i = 0; i < values.size(); i++) {
+        path.addLast(RegoInt32.of(i));
+        collect(path, values.get(i), pairs);
+        path.removeLast();
+      }
+    } else if (value instanceof RegoSet) {
+      // Sets are keyed by their own members, so a member's path segment is the member itself.
+      for (RegoValue member : ((RegoSet) value).getValue()) {
+        path.addLast(member);
+        collect(path, member, pairs);
+        path.removeLast();
+      }
+    }
+  }
+}

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

@@ -0,0 +1,114 @@
+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.assertTrue;
+
+import io.github.open_policy_agent.opa.ast.types.RegoArray;
+import io.github.open_policy_agent.opa.ast.types.RegoInt32;
+import io.github.open_policy_agent.opa.ast.types.RegoNull;
+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.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Expected pairs are the output of {@code opa eval} on the equivalent {@code walk(x, [p, v])} query
+ * for the same input, so these assertions pin Go-OPA parity rather than just internal consistency.
+ */
+public class GraphBuiltinsTest {
+
+  private final GraphBuiltins graphBuiltins = new GraphBuiltins();
+  private final EvaluationContext ctx = new EvaluationContext.Builder().build();
+
+  private Set<String> walkPairs(RegoValue input) {
+    RegoValue result = graphBuiltins.walk(ctx, new RegoValue[] {input});
+    assertTrue(result instanceof RegoSet, "walk must return a set of [path, value] pairs");
+    return ((RegoSet) result)
+        .getValue().stream().map(RegoValue::toString).collect(Collectors.toSet());
+  }
+
+  @Test
+  public void testWalkNestedObjectAndArray() {
+    // {"a": {"b": [1, {"c": "d"}]}, "e": null}
+    RegoObject inner = new RegoObject(Map.of(new RegoString("c"), new RegoString("d")));
+    RegoArray b = new RegoArray(List.of(RegoInt32.of(1), inner));
+    Map<RegoValue, RegoValue> root = new LinkedHashMap<>();
+    root.put(new RegoString("a"), new RegoObject(Map.of(new RegoString("b"), b)));
+    root.put(new RegoString("e"), RegoNull.INSTANCE);
+    RegoObject input = new RegoObject(root);
+
+    Set<String> pairs = walkPairs(input);
+
+    // The root is always emitted with an empty path, plus one pair per nested document.
+    assertEquals(7, pairs.size());
+    assertTrue(pairs.contains("[[], " + input + "]"), "root pair with empty path: " + pairs);
+    assertTrue(pairs.contains("[[\"a\", \"b\", 0], 1]"), "array index in path: " + pairs);
+    assertTrue(pairs.contains("[[\"a\", \"b\", 1, \"c\"], \"d\"]"), "leaf under array: " + pairs);
+    assertTrue(
+        pairs.contains("[[\"a\", \"b\", 1], " + inner + "]"), "array element itself: " + pairs);
+    assertTrue(pairs.contains("[[\"a\", \"b\"], " + b + "]"), "array itself: " + pairs);
+    assertTrue(pairs.contains("[[\"e\"], null]"), "null leaf is walked: " + pairs);
+  }
+
+  @Test
+  public void testWalkDescendsIntoSetsKeyedByMember() {
+    // walk({"k": {"x", "y"}, "n": 3} -- a set member's path segment is the member itself.
+    RegoSet members = new RegoSet(false);
+    members.addValue(new RegoString("x"));
+    members.addValue(new RegoString("y"));
+    Map<RegoValue, RegoValue> root = new LinkedHashMap<>();
+    root.put(new RegoString("k"), members);
+    root.put(new RegoString("n"), RegoInt32.of(3));
+
+    Set<String> pairs = walkPairs(new RegoObject(root));
+
+    assertEquals(5, pairs.size());
+    assertTrue(pairs.contains("[[\"k\", \"x\"], \"x\"]"), "set member keyed by itself: " + pairs);
+    assertTrue(pairs.contains("[[\"k\", \"y\"], \"y\"]"), "set member keyed by itself: " + pairs);
+    assertTrue(pairs.contains("[[\"n\"], 3]"), "sibling scalar: " + pairs);
+  }
+
+  @Test
+  public void testWalkObjectWithCompositeKeyUsesKeyAsPathSegment() {
+    // walk({["a", 1]: "v"}) -- an object key can itself be composite, so a path segment can be an
+    // array. opa eval yields the root plus a pair whose path is [["a", 1]].
+    RegoArray compositeKey = new RegoArray(List.of(new RegoString("a"), RegoInt32.of(1)));
+    RegoObject input = new RegoObject(Map.of(compositeKey, new RegoString("v")));
+
+    Set<String> pairs = walkPairs(input);
+
+    assertEquals(2, pairs.size());
+    assertTrue(pairs.contains("[[], " + input + "]"), "root pair: " + pairs);
+    assertTrue(
+        pairs.contains("[[" + compositeKey + "], \"v\"]"),
+        "composite key becomes one path segment: " + pairs);
+  }
+
+  @Test
+  public void testWalkScalarYieldsOnlyRoot() {
+    Set<String> pairs = walkPairs(RegoInt32.of(5));
+
+    assertEquals(Set.of("[[], 5]"), pairs);
+  }
+
+  @Test
+  public void testWalkEmptyCollectionYieldsOnlyRoot() {
+    RegoObject empty = new RegoObject();
+
+    assertEquals(Set.of("[[], " + empty + "]"), walkPairs(empty));
+  }
+
+  @Test
+  public void testWalkUndefinedArgumentYieldsNoPairs() {
+    // An undefined operand makes the expression undefined in Go-OPA; the plan scans the returned
+    // collection, so an empty set is what produces zero iterations here.
+    assertEquals(Set.of(), walkPairs(null));
+  }
+}

opa-evaluator/src/test/resources/compliance/known-missing-builtins.txt

@@ -112,6 +112,3 @@ urlquery.encode_object
 
 uuid.parse
 uuid.rfc4122
-
-# Delete this line along with the walk implementation in #141.
-walk

opa-jackson/src/main/java/io/github/open_policy_agent/opa/jackson/RegoValueModule.java

@@ -1,5 +1,6 @@
 package io.github.open_policy_agent.opa.jackson;
 
+import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.DeserializationContext;
@@ -22,6 +23,7 @@
 import io.github.open_policy_agent.opa.ast.types.RegoValue;
 
 import java.io.IOException;
+import java.io.StringWriter;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.util.Iterator;
@@ -128,12 +130,15 @@ public void serialize(RegoSet v, JsonGenerator g, SerializerProvider p) throws I
   }
 
   private static final class RegoObjectSerializer extends JsonSerializer<RegoObject> {
+    /** Used to render composite object keys on their own; kept static since factories are heavy. */
+    private static final JsonFactory KEY_FACTORY = new JsonFactory();
+
     @Override
     public void serialize(RegoObject v, JsonGenerator g, SerializerProvider p) throws IOException {
       // OPA (Go) emits sorted keys; preserve that for round-trip fidelity.
       Map<String, RegoValue> sorted = new TreeMap<>();
       for (Map.Entry<RegoValue, RegoValue> entry : v.getProperties().entrySet()) {
-        sorted.put(keyToString(entry.getKey()), entry.getValue());
+        sorted.put(keyToString(entry.getKey(), p), entry.getValue());
       }
       g.writeStartObject();
       for (Map.Entry<String, RegoValue> entry : sorted.entrySet()) {
@@ -143,14 +148,22 @@ public void serialize(RegoObject v, JsonGenerator g, SerializerProvider p) throw
       g.writeEndObject();
     }
 
-    private static String keyToString(RegoValue key) {
+    private static String keyToString(RegoValue key, SerializerProvider p) throws IOException {
       if (key instanceof RegoString) {
         return ((RegoString) key).getValue();
       }
       if (key instanceof RegoNumber) {
         return ((RegoNumber) key).getBigIntValue().toString();
       }
-      return key.toString();
+      // Composite keys (arrays/sets/objects, e.g. the paths produced by walk) are named by their
+      // compact JSON encoding, matching Go-OPA. Re-serializing through the provider keeps nested
+      // values consistent with the rest of the document; RegoValue.toString() would instead leak
+      // Java's collection formatting ("[a, 0]" rather than ["a",0]).
+      StringWriter out = new StringWriter();
+      try (JsonGenerator keyGen = KEY_FACTORY.createGenerator(out)) {
+        p.defaultSerializeValue(key, keyGen);
+      }
+      return out.toString();
     }
   }
 

opa-jackson/src/test/java/io/github/open_policy_agent/opa/jackson/RegoValueModuleTest.java

@@ -94,6 +94,77 @@ void serialize_regoObject_numericKeysCoercedToString() throws IOException {
     assertThat(mapper.writeValueAsString(obj)).isEqualTo("{\"7\":\"seven\"}");
   }
 
+  @Test
+  void serialize_regoObject_compositeKeysUseCompactJson() throws IOException {
+    // Composite keys arise from e.g. `{path: value | [path, value] := walk(x)}`. OPA names them by
+    // their compact JSON encoding, so no spaces after the separators.
+    RegoObject obj = new RegoObject();
+    RegoArray path = new RegoArray();
+    path.addValue(new RegoString("a"));
+    path.addValue(RegoInt32.of(0));
+    path.addValue(new RegoString("b"));
+    obj.setProp(path, new RegoString("AA"));
+
+    assertThat(mapper.writeValueAsString(obj)).isEqualTo("{\"[\\\"a\\\",0,\\\"b\\\"]\":\"AA\"}");
+  }
+
+  @Test
+  void serialize_regoObject_compositeKeysOfEveryKind() throws IOException {
+    // Sets and objects can be keys too, not just arrays. Expected output is the `opa eval` result
+    // for {{"s1", "s2"}: "set-key", {"k": [1, 2]}: "obj-key", [["nested"], 3]: "nested-key"}.
+    RegoSet setKey = new RegoSet(false);
+    setKey.addValue(new RegoString("s1"));
+    setKey.addValue(new RegoString("s2"));
+
+    RegoArray inner = new RegoArray();
+    inner.addValue(RegoInt32.of(1));
+    inner.addValue(RegoInt32.of(2));
+    RegoObject objKey = new RegoObject();
+    objKey.setProp(new RegoString("k"), inner);
+
+    RegoArray nestedKey = new RegoArray();
+    RegoArray nestedFirst = new RegoArray();
+    nestedFirst.addValue(new RegoString("nested"));
+    nestedKey.addValue(nestedFirst);
+    nestedKey.addValue(RegoInt32.of(3));
+
+    RegoObject obj = new RegoObject();
+    obj.setProp(objKey, new RegoString("obj-key"));
+    obj.setProp(setKey, new RegoString("set-key"));
+    obj.setProp(nestedKey, new RegoString("nested-key"));
+
+    // Keys sort lexicographically by their JSON text, which is how OPA orders them: " < [ < {.
+    assertThat(mapper.writeValueAsString(obj))
+        .isEqualTo(
+            "{\"[\\\"s1\\\",\\\"s2\\\"]\":\"set-key\","
+                + "\"[[\\\"nested\\\"],3]\":\"nested-key\","
+                + "\"{\\\"k\\\":[1,2]}\":\"obj-key\"}");
+  }
+
+  @Test
+  void serialize_regoObject_compositeKeyOrderFollowsJsonText() throws IOException {
+    // Sorting on the compact encoding is what makes this match OPA: with the old spaced form,
+    // "[\"a\", 10]" sorted before "[\"a\", \"z\"]" (space < quote), i.e. the wrong order.
+    RegoObject obj = new RegoObject();
+    obj.setProp(arrayKey(new RegoString("a"), RegoInt32.of(2)), new RegoString("two"));
+    obj.setProp(arrayKey(new RegoString("a"), RegoInt32.of(10)), new RegoString("ten"));
+    obj.setProp(arrayKey(new RegoString("a"), new RegoString("z")), new RegoString("zed"));
+
+    assertThat(mapper.writeValueAsString(obj))
+        .isEqualTo(
+            "{\"[\\\"a\\\",\\\"z\\\"]\":\"zed\","
+                + "\"[\\\"a\\\",10]\":\"ten\","
+                + "\"[\\\"a\\\",2]\":\"two\"}");
+  }
+
+  private static RegoArray arrayKey(RegoValue... items) {
+    RegoArray key = new RegoArray();
+    for (RegoValue item : items) {
+      key.addValue(item);
+    }
+    return key;
+  }
+
   @Test
   void serialize_nestedStructure() throws IOException {
     RegoObject obj = new RegoObject();