[ggj][nestedsig] feat: add DFS traversal to Trie

src/main/java/com/google/api/generator/util/Trie.java

@@ -14,17 +14,20 @@
 
 package com.google.api.generator.util;
 
-import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 
 /**
  * A common-prefix trie. T represents the type of each "char" in a word (which is a T-typed list).
  */
 public class Trie<T> {
   private class Node<T> {
     final T chr;
-    Map<T, Node> children = new HashMap<>();
+    // Maintain insertion order to enable deterministic test output.
+    Map<T, Node<T>> children = new LinkedHashMap<>();
     boolean isLeaf;
 
     Node() {
@@ -43,7 +46,7 @@ public Trie() {
   }
 
   public void insert(List<T> word) {
-    Map<T, Node> children = root.children;
+    Map<T, Node<T>> children = root.children;
     for (int i = 0; i < word.size(); i++) {
       T chr = word.get(i);
       Node t;
@@ -71,8 +74,53 @@ public boolean hasPrefix(List<T> prefix) {
     return searchNode(prefix) != null;
   }
 
+  /**
+   * Reduces the trie to a single value, via a DFS traversal.
+   *
+   * @param parentPreprocFn Transforms a parent node into an R-typed base value for consumption by
+   *     the child nodes. The rest of the children will compute their values using this as a base as
+   *     well, so it accumulates computational results as the traversal progresses. Does not handle
+   *     the root node (i.e. when {@code chr} is null).
+   * @param leafReduceFn Transforms a child node into an R-typed value using the value computed by
+   *     the parent nodes' preprocessing functions.
+   * @param parentPostprocFn Transforms the post-traversal result (from the child nodes) into
+   *     R-typed values, further building upon {@code baseValue}. Must handle the root node, i.e.
+   *     when {@code chr} is null.
+   * @param baseValue The base value upon which subsequent reductions will be performed. Ensure this
+   *     is a type that can accumulate values, such as StringBuilder. An immutable type such as
+   *     String will not work here.
+   */
+  public <R> R dfsTraverseAndReduce(
+      Function<T, R> parentPreprocFn,
+      TriFunction<T, R, R, R> parentPostprocFn,
+      BiFunction<T, R, R> leafReduceFn,
+      R baseValue) {
+    return dfsTraverseAndReduce(root, parentPreprocFn, parentPostprocFn, leafReduceFn, baseValue);
+  }
+
+  /** Traverses the trie DFS-style, reducing all values into a single one on {@code baseValue}. */
+  private <R> R dfsTraverseAndReduce(
+      Node<T> node,
+      Function<T, R> parentPreprocFn,
+      TriFunction<T, R, R, R> parentPostprocFn,
+      BiFunction<T, R, R> leafReduceFn,
+      R baseValue) {
+    if (node.isLeaf) {
+      return leafReduceFn.apply(node.chr, baseValue);
+    }
+
+    R leafReducedValue = node.chr == null ? baseValue : parentPreprocFn.apply(node.chr);
+    for (Map.Entry<T, Node<T>> e : node.children.entrySet()) {
+      // Thread the parent value through each of the children, and accumulate it.
+      leafReducedValue =
+          dfsTraverseAndReduce(
+              e.getValue(), parentPreprocFn, parentPostprocFn, leafReduceFn, leafReducedValue);
+    }
+    return parentPostprocFn.apply(node.chr, baseValue, leafReducedValue);
+  }
+
   private Node searchNode(List<T> word) {
-    Map<T, Node> children = root.children;
+    Map<T, Node<T>> children = root.children;
     Node t = null;
     for (int i = 0; i < word.size(); i++) {
       T chr = word.get(i);

src/test/java/com/google/api/generator/util/BUILD.bazel

@@ -17,6 +17,7 @@ filegroup(
     test_class = "com.google.api.generator.util.{0}".format(test_name),
     deps = [
         "//src/main/java/com/google/api/generator/util",
+        "//src/test/java/com/google/api/generator/testutils",
         "@com_google_guava_guava//jar",
         "@com_google_truth_truth//jar",
         "@junit_junit//jar",

src/test/java/com/google/api/generator/util/TrieTest.java

@@ -14,17 +14,20 @@
 
 package com.google.api.generator.util;
 
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
+import com.google.api.generator.testutils.LineFormatter;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import org.junit.Test;
 
 public class TrieTest {
   @Test
-  public void stringTrie() {
+  public void insertAndSearch_stringTrie() {
     Trie<String> trie = new Trie<>();
 
     Function<String, List<String>> wordToCharListFn = w -> Arrays.asList(w.split("(?!^)"));
@@ -43,7 +46,7 @@ public void stringTrie() {
   }
 
   @Test
-  public void multiStringTrie() {
+  public void insertAndSearch_multiStringTrie() {
     Trie<String> trie = new Trie<>();
     assertFalse(trie.search(Arrays.asList("user", "identity", "name")));
 
@@ -62,4 +65,191 @@ public void multiStringTrie() {
     assertFalse(trie.hasPrefix(Arrays.asList("identity")));
     assertFalse(trie.hasPrefix(Arrays.asList("contact")));
   }
+
+  @Test
+  public void dfsTraverseAndReduce_emptyTrie() {
+    // Add up points in the tree, where each parent gets (num child node points) * 2 + 1.
+    int baseValue = 0;
+    Function<String, Integer> parentPreprocFn = nodeVal -> new Integer(0);
+    BiFunction<String, Integer, Integer> leafReduceFn =
+        (nodeVal, accVal) -> new Integer(accVal + 1);
+    TriFunction<String, Integer, Integer, Integer> parentPostprocFn =
+        (nodeVal, baseVal, accVal) -> new Integer(nodeVal == null ? 0 : accVal * 2 + 1);
+
+    Trie<String> trie = new Trie<>();
+    int finalValue = trie.dfsTraverseAndReduce(parentPreprocFn, parentPostprocFn, leafReduceFn, 0);
+    assertEquals(0, finalValue);
+  }
+
+  @Test
+  public void dfsTraverseAndReduce_singleNodeTrie() {
+    // Add up points in the tree, where each parent gets (num child node points) * 2 + 1.
+    int baseValue = 0;
+    Function<String, Integer> parentPreprocFn = nodeVal -> new Integer(0);
+    BiFunction<String, Integer, Integer> leafReduceFn =
+        (nodeVal, accVal) -> new Integer(accVal + 1);
+    TriFunction<String, Integer, Integer, Integer> parentPostprocFn =
+        (nodeVal, baseVal, accVal) -> new Integer(nodeVal == null ? accVal : accVal * 2 + 1);
+
+    Trie<String> trie = new Trie<>();
+    trie.insert(Arrays.asList("user"));
+    int finalValue = trie.dfsTraverseAndReduce(parentPreprocFn, parentPostprocFn, leafReduceFn, 0);
+    assertEquals(1, finalValue);
+  }
+
+  @Test
+  public void dfsTraverseAndReduce_oneParentOneChildBranchTrie() {
+    Function<String, String> toUpperCaseFn = s -> s.substring(0, 1).toUpperCase() + s.substring(1);
+    Function<String, StringBuilder> parentPreprocFn =
+        nodeVal ->
+            new StringBuilder(String.format("%s.newBuilder()", toUpperCaseFn.apply(nodeVal)));
+    BiFunction<String, StringBuilder, StringBuilder> leafReduceFn =
+        (nodeVal, parentAccVal) -> {
+          parentAccVal.append(
+              String.format(
+                  ".set%s(\"%s\")", toUpperCaseFn.apply(nodeVal.toString()), nodeVal.toString()));
+          return parentAccVal;
+        };
+    TriFunction<String, StringBuilder, StringBuilder, StringBuilder> parentPostprocFn =
+        (nodeVal, baseVal, accVal) -> {
+          boolean isRootNode = nodeVal == null;
+          if (!isRootNode) {
+            baseVal.append(
+                String.format(
+                    ".set%s(%s.build())",
+                    toUpperCaseFn.apply(nodeVal.toString()), accVal.toString()));
+          }
+          return isRootNode ? accVal : baseVal;
+        };
+    StringBuilder baseVal = new StringBuilder("RequestType.newBuilder()");
+
+    Trie<String> trie = new Trie<>();
+    trie.insert(Arrays.asList("user", "identity"));
+    String finalValue =
+        trie.dfsTraverseAndReduce(parentPreprocFn, parentPostprocFn, leafReduceFn, baseVal)
+            .toString();
+
+    assertEquals(
+        LineFormatter.lines(
+            "RequestType.newBuilder()",
+            ".setUser(User.newBuilder()",
+            ".setIdentity(\"identity\").build())",
+            ".build();"),
+        finalValue + ".build();");
+  }
+
+  @Test
+  public void dfsTraverseAndReduce_oneDeepBranchTrie() {
+    // Add up points in the tree, where each parent gets (num child node points) * 2 + 1.
+    int simpleBaseValue = 0;
+    Function<String, Integer> simpleParentPreprocFn = nodeVal -> new Integer(0);
+    BiFunction<String, Integer, Integer> simpleLeafReduceFn =
+        (nodeVal, accVal) -> new Integer(accVal + 1);
+    TriFunction<String, Integer, Integer, Integer> simpleParentPostprocFn =
+        (nodeVal, baseVal, accVal) -> new Integer(nodeVal == null ? accVal : accVal * 2 + 1);
+
+    Trie<String> trie = new Trie<>();
+    trie.insert(Arrays.asList("user", "identity", "name", "firstName"));
+    int simpleFinalValue =
+        trie.dfsTraverseAndReduce(
+            simpleParentPreprocFn, simpleParentPostprocFn, simpleLeafReduceFn, simpleBaseValue);
+    assertEquals(15, simpleFinalValue);
+
+    Function<String, String> toUpperCaseFn = s -> s.substring(0, 1).toUpperCase() + s.substring(1);
+    Function<String, StringBuilder> parentPreprocFn =
+        nodeVal ->
+            new StringBuilder(String.format("%s.newBuilder()", toUpperCaseFn.apply(nodeVal)));
+    BiFunction<String, StringBuilder, StringBuilder> leafReduceFn =
+        (nodeVal, parentAccVal) -> {
+          parentAccVal.append(
+              String.format(
+                  ".set%s(\"%s\")", toUpperCaseFn.apply(nodeVal.toString()), nodeVal.toString()));
+          return parentAccVal;
+        };
+    TriFunction<String, StringBuilder, StringBuilder, StringBuilder> parentPostprocFn =
+        (nodeVal, baseVal, accVal) -> {
+          boolean isRootNode = nodeVal == null;
+          if (!isRootNode) {
+            baseVal.append(
+                String.format(
+                    ".set%s(%s.build())",
+                    toUpperCaseFn.apply(nodeVal.toString()), accVal.toString()));
+          }
+          return isRootNode ? accVal : baseVal;
+        };
+    StringBuilder baseVal = new StringBuilder("RequestType.newBuilder()");
+
+    String finalValue =
+        trie.dfsTraverseAndReduce(parentPreprocFn, parentPostprocFn, leafReduceFn, baseVal)
+            .toString();
+
+    assertEquals(
+        LineFormatter.lines(
+            "RequestType.newBuilder()",
+            ".setUser(User.newBuilder()",
+            ".setIdentity(Identity.newBuilder()",
+            ".setName(Name.newBuilder()",
+            ".setFirstName(\"firstName\").build())",
+            ".build())",
+            ".build())",
+            ".build();"),
+        finalValue + ".build();");
+  }
+
+  @Test
+  public void dfsTraverseAndReduce_depthAndBreathTrie() {
+    Function<String, String> toUpperCaseFn = s -> s.substring(0, 1).toUpperCase() + s.substring(1);
+    Function<String, StringBuilder> parentPreprocFn =
+        nodeVal ->
+            new StringBuilder(String.format("%s.newBuilder()", toUpperCaseFn.apply(nodeVal)));
+    BiFunction<String, StringBuilder, StringBuilder> leafReduceFn =
+        (nodeVal, parentAccVal) -> {
+          parentAccVal.append(
+              String.format(
+                  ".set%s(\"%s\")", toUpperCaseFn.apply(nodeVal.toString()), nodeVal.toString()));
+          return parentAccVal;
+        };
+    TriFunction<String, StringBuilder, StringBuilder, StringBuilder> parentPostprocFn =
+        (nodeVal, baseVal, accVal) -> {
+          boolean isRootNode = nodeVal == null;
+          if (!isRootNode) {
+            baseVal.append(
+                String.format(
+                    ".set%s(%s.build())",
+                    toUpperCaseFn.apply(nodeVal.toString()), accVal.toString()));
+          }
+          return isRootNode ? accVal : baseVal;
+        };
+    StringBuilder baseVal = new StringBuilder("RequestType.newBuilder()");
+
+    Trie<String> trie = new Trie<>();
+    trie.insert(Arrays.asList("user", "identity", "name", "firstName"));
+    trie.insert(Arrays.asList("user", "identity", "name", "lastName"));
+    trie.insert(Arrays.asList("user", "email"));
+    trie.insert(Arrays.asList("user", "age"));
+    trie.insert(Arrays.asList("user", "hobby", "hobbyName"));
+    trie.insert(Arrays.asList("user", "hobby", "frequency"));
+
+    String finalValue =
+        trie.dfsTraverseAndReduce(parentPreprocFn, parentPostprocFn, leafReduceFn, baseVal)
+            .toString();
+
+    assertEquals(
+        LineFormatter.lines(
+            "RequestType.newBuilder()",
+            ".setUser(User.newBuilder()",
+            ".setIdentity(Identity.newBuilder()",
+            ".setName(Name.newBuilder()",
+            ".setFirstName(\"firstName\")",
+            ".setLastName(\"lastName\").build())", // Name.newBuilder().build().
+            ".build())", // Identity.newBuilder().build().
+            ".setEmail(\"email\")",
+            ".setAge(\"age\")",
+            ".setHobby(Hobby.newBuilder()",
+            ".setHobbyName(\"hobbyName\")",
+            ".setFrequency(\"frequency\").build())", // Hobby.newBuilder().build().
+            ".build())", // User.newBuilder().build().
+            ".build();"),
+        finalValue + ".build();");
+  }
 }