Refactor PathTrie to tidy it up (#107542)

server/src/main/java/org/elasticsearch/common/path/PathTrie.java

@@ -10,13 +10,15 @@ package org.elasticsearch.common.path;
 
 import org.elasticsearch.common.collect.Iterators;
 
+import java.util.Arrays;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
-import java.util.NoSuchElementException;
-import java.util.function.BiFunction;
+import java.util.function.BinaryOperator;
 import java.util.function.Supplier;
+import java.util.function.UnaryOperator;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.unmodifiableMap;
@@ -50,52 +52,43 @@ public class PathTrie<T> {
         TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED
     );
 
-    public interface Decoder {
-        String decode(String value);
-    }
-
-    private final Decoder decoder;
+    private final UnaryOperator<String> decoder;
     private final TrieNode root;
     private T rootValue;
 
     private static final String SEPARATOR = "/";
     private static final String WILDCARD = "*";
 
-    public PathTrie(Decoder decoder) {
+    public PathTrie(UnaryOperator<String> decoder) {
         this.decoder = decoder;
-        root = new TrieNode(SEPARATOR, null, WILDCARD);
+        root = new TrieNode(SEPARATOR, null);
     }
 
-    public class TrieNode {
-        private transient String key;
-        private transient T value;
-        private final String wildcard;
-
-        private transient String namedWildcard;
-
+    private class TrieNode {
+        private T value;
+        private String namedWildcard;
         private Map<String, TrieNode> children;
 
-        private TrieNode(String key, T value, String wildcard) {
-            this.key = key;
-            this.wildcard = wildcard;
+        private TrieNode(String key, T value) {
             this.value = value;
             this.children = emptyMap();
             if (isNamedWildcard(key)) {
-                namedWildcard = key.substring(key.indexOf('{') + 1, key.indexOf('}'));
+                updateNamedWildcard(key);
             } else {
                 namedWildcard = null;
             }
         }
 
-        private void updateKeyWithNamedWildcard(String key) {
-            this.key = key;
-            String newNamedWildcard = key.substring(key.indexOf('{') + 1, key.indexOf('}'));
-            if (namedWildcard != null && newNamedWildcard.equals(namedWildcard) == false) {
-                throw new IllegalArgumentException(
-                    "Trying to use conflicting wildcard names for same path: " + namedWildcard + " and " + newNamedWildcard
-                );
+        private void updateNamedWildcard(String key) {
+            String newNamedWildcard = key.substring(1, key.length() - 1);
+            if (newNamedWildcard.equals(namedWildcard) == false) {
+                if (namedWildcard != null) {
+                    throw new IllegalArgumentException(
+                        "Trying to use conflicting wildcard names for same path: " + namedWildcard + " and " + newNamedWildcard
+                    );
+                }
+                namedWildcard = newNamedWildcard;
             }
-            namedWildcard = newNamedWildcard;
         }
 
         private void addInnerChild(String key, TrieNode child) {
@@ -110,16 +103,17 @@ public class PathTrie<T> {
             String token = path[index];
             String key = token;
             if (isNamedWildcard(token)) {
-                key = wildcard;
+                key = WILDCARD;
             }
+
             TrieNode node = children.get(key);
             if (node == null) {
                 T nodeValue = index == path.length - 1 ? value : null;
-                node = new TrieNode(token, nodeValue, wildcard);
+                node = new TrieNode(token, nodeValue);
                 addInnerChild(key, node);
             } else {
                 if (isNamedWildcard(token)) {
-                    node.updateKeyWithNamedWildcard(token);
+                    node.updateNamedWildcard(token);
                 }
                 /*
                  * If the target node already exists, but is without a value,
@@ -139,22 +133,23 @@ public class PathTrie<T> {
             node.insert(path, index + 1, value);
         }
 
-        private synchronized void insertOrUpdate(String[] path, int index, T value, BiFunction<T, T, T> updater) {
+        private synchronized void insertOrUpdate(String[] path, int index, T value, BinaryOperator<T> updater) {
             if (index >= path.length) return;
 
             String token = path[index];
             String key = token;
             if (isNamedWildcard(token)) {
-                key = wildcard;
+                key = WILDCARD;
             }
+
             TrieNode node = children.get(key);
             if (node == null) {
                 T nodeValue = index == path.length - 1 ? value : null;
-                node = new TrieNode(token, nodeValue, wildcard);
+                node = new TrieNode(token, nodeValue);
                 addInnerChild(key, node);
             } else {
                 if (isNamedWildcard(token)) {
-                    node.updateKeyWithNamedWildcard(token);
+                    node.updateNamedWildcard(token);
                 }
                 /*
                  * If the target node already exists, but is without a value,
@@ -173,7 +168,7 @@ public class PathTrie<T> {
         }
 
         private static boolean isNamedWildcard(String key) {
-            return key.indexOf('{') != -1 && key.indexOf('}') != -1;
+            return key.charAt(0) == '{' && key.charAt(key.length() - 1) == '}';
         }
 
         private String namedWildcard() {
@@ -184,7 +179,7 @@ public class PathTrie<T> {
             return namedWildcard != null;
         }
 
-        public T retrieve(String[] path, int index, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
+        private T retrieve(String[] path, int index, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
             if (index >= path.length) return null;
 
             String token = path[index];
@@ -193,7 +188,7 @@ public class PathTrie<T> {
 
             if (node == null) {
                 if (trieMatchingMode == TrieMatchingMode.WILDCARD_NODES_ALLOWED) {
-                    node = children.get(wildcard);
+                    node = children.get(WILDCARD);
                     if (node == null) {
                         return null;
                     }
@@ -202,7 +197,7 @@ public class PathTrie<T> {
                     /*
                      * Allow root node wildcard matches.
                      */
-                    node = children.get(wildcard);
+                    node = children.get(WILDCARD);
                     if (node == null) {
                         return null;
                     }
@@ -211,7 +206,7 @@ public class PathTrie<T> {
                     /*
                      * Allow leaf node wildcard matches.
                      */
-                    node = children.get(wildcard);
+                    node = children.get(WILDCARD);
                     if (node == null) {
                         return null;
                     }
@@ -220,32 +215,33 @@ public class PathTrie<T> {
                     return null;
                 }
             } else {
+                TrieNode wildcardNode;
                 if (index + 1 == path.length
                     && node.value == null
-                    && children.get(wildcard) != null
-                    && EXPLICIT_OR_ROOT_WILDCARD.contains(trieMatchingMode) == false) {
+                    && EXPLICIT_OR_ROOT_WILDCARD.contains(trieMatchingMode) == false
+                    && (wildcardNode = children.get(WILDCARD)) != null) {
                     /*
                      * If we are at the end of the path, the current node does not have a value but
                      * there is a child wildcard node, use the child wildcard node.
                      */
-                    node = children.get(wildcard);
+                    node = wildcardNode;
                     usedWildcard = true;
                 } else if (index == 1
                     && node.value == null
-                    && children.get(wildcard) != null
-                    && trieMatchingMode == TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED) {
+                    && trieMatchingMode == TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED
+                    && (wildcardNode = children.get(WILDCARD)) != null) {
                         /*
                          * If we are at the root, and root wildcards are allowed, use the child wildcard
                          * node.
                          */
-                        node = children.get(wildcard);
+                        node = wildcardNode;
                         usedWildcard = true;
                     } else {
-                        usedWildcard = token.equals(wildcard);
+                        usedWildcard = token.equals(WILDCARD);
                     }
             }
 
-            put(params, node, token);
+            recordWildcardParam(params, node, token);
 
             if (index == (path.length - 1)) {
                 return node.value;
@@ -253,9 +249,9 @@ public class PathTrie<T> {
 
             T nodeValue = node.retrieve(path, index + 1, params, trieMatchingMode);
             if (nodeValue == null && usedWildcard == false && trieMatchingMode != TrieMatchingMode.EXPLICIT_NODES_ONLY) {
-                node = children.get(wildcard);
+                node = children.get(WILDCARD);
                 if (node != null) {
-                    put(params, node, token);
+                    recordWildcardParam(params, node, token);
                     nodeValue = node.retrieve(path, index + 1, params, trieMatchingMode);
                 }
             }
@@ -263,13 +259,13 @@ public class PathTrie<T> {
             return nodeValue;
         }
 
-        private void put(Map<String, String> params, TrieNode node, String value) {
+        private void recordWildcardParam(Map<String, String> params, TrieNode node, String value) {
             if (params != null && node.isNamedWildcard()) {
-                params.put(node.namedWildcard(), decoder.decode(value));
+                params.put(node.namedWildcard(), decoder.apply(value));
             }
         }
 
-        Iterator<T> allNodeValues() {
+        private Iterator<T> allNodeValues() {
             final Iterator<T> childrenIterator = Iterators.flatMap(children.values().iterator(), TrieNode::allNodeValues);
             if (value == null) {
                 return childrenIterator;
@@ -277,11 +273,6 @@ public class PathTrie<T> {
                 return Iterators.concat(Iterators.single(value), childrenIterator);
             }
         }
-
-        @Override
-        public String toString() {
-            return key;
-        }
     }
 
     public void insert(String path, T value) {
@@ -308,7 +299,7 @@ public class PathTrie<T> {
      * </pre>
      * allowing the value to be updated if desired.
      */
-    public void insertOrUpdate(String path, T value, BiFunction<T, T, T> updater) {
+    public void insertOrUpdate(String path, T value, BinaryOperator<T> updater) {
         String[] strings = path.split(SEPARATOR);
         if (strings.length == 0) {
             if (rootValue != null) {
@@ -334,8 +325,8 @@ public class PathTrie<T> {
         return retrieve(path, params, TrieMatchingMode.WILDCARD_NODES_ALLOWED);
     }
 
-    public T retrieve(String path, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
-        if (path.length() == 0) {
+    T retrieve(String path, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
+        if (path.isEmpty()) {
             return rootValue;
         }
         String[] strings = path.split(SEPARATOR);
@@ -353,29 +344,12 @@ public class PathTrie<T> {
     }
 
     /**
-     * Returns an iterator of the objects stored in the {@code PathTrie}, using
+     * Returns a stream of the objects stored in the {@code PathTrie}, using
      * all possible {@code TrieMatchingMode} modes. The {@code paramSupplier}
-     * is called between each invocation of {@code next()} to supply a new map
-     * of parameters.
+     * is called for each mode to supply a new map of parameters.
      */
-    public Iterator<T> retrieveAll(String path, Supplier<Map<String, String>> paramSupplier) {
-        return new Iterator<>() {
-
-            private int mode;
-
-            @Override
-            public boolean hasNext() {
-                return mode < TrieMatchingMode.values().length;
-            }
-
-            @Override
-            public T next() {
-                if (hasNext() == false) {
-                    throw new NoSuchElementException("called next() without validating hasNext()! no more modes available");
-                }
-                return retrieve(path, paramSupplier.get(), TrieMatchingMode.values()[mode++]);
-            }
-        };
+    public Stream<T> retrieveAll(String path, Supplier<Map<String, String>> paramSupplier) {
+        return Arrays.stream(TrieMatchingMode.values()).map(m -> retrieve(path, paramSupplier.get(), m));
     }
 
     public Iterator<T> allNodeValues() {

server/src/main/java/org/elasticsearch/rest/RestController.java

@@ -644,7 +644,7 @@ public class RestController implements HttpServerTransport.Dispatcher {
         // we use rawPath since we don't want to decode it while processing the path resolution
         // so we can handle things like:
         // my_index/my_type/http%3A%2F%2Fwww.google.com
-        return handlers.retrieveAll(rawPath, paramsSupplier);
+        return handlers.retrieveAll(rawPath, paramsSupplier).iterator();
     }
 
     /**

server/src/main/java/org/elasticsearch/rest/RestUtils.java

@@ -9,7 +9,6 @@
 package org.elasticsearch.rest;
 
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.path.PathTrie;
 import org.elasticsearch.core.Booleans;
 
 import java.nio.charset.Charset;
@@ -17,6 +16,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.UnaryOperator;
 import java.util.regex.Pattern;
 
 import static org.elasticsearch.rest.RestRequest.PATH_RESTRICTED;
@@ -28,7 +28,7 @@ public class RestUtils {
      */
     private static final boolean DECODE_PLUS_AS_SPACE = Booleans.parseBoolean(System.getProperty("es.rest.url_plus_as_space", "false"));
 
-    public static final PathTrie.Decoder REST_DECODER = RestUtils::decodeComponent;
+    public static final UnaryOperator<String> REST_DECODER = RestUtils::decodeComponent;
 
     public static void decodeQueryString(String s, int fromIndex, Map<String, String> params) {
         if (fromIndex < 0) {

server/src/test/java/org/elasticsearch/common/path/PathTrieTests.java

@@ -13,20 +13,16 @@ import org.elasticsearch.rest.RestUtils;
 import org.elasticsearch.test.ESTestCase;
 
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.Map;
+import java.util.function.UnaryOperator;
 
+import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.nullValue;
 
 public class PathTrieTests extends ESTestCase {
 
-    public static final PathTrie.Decoder NO_DECODER = new PathTrie.Decoder() {
-        @Override
-        public String decode(String value) {
-            return value;
-        }
-    };
+    public static final UnaryOperator<String> NO_DECODER = UnaryOperator.identity();
 
     public void testPath() {
         PathTrie<String> trie = new PathTrie<>(NO_DECODER);
@@ -50,9 +46,7 @@ public class PathTrieTests extends ESTestCase {
 
         Map<String, String> params = new HashMap<>();
         assertThat(trie.retrieve("index1/insert/12", params), equalTo("bingo"));
-        assertThat(params.size(), equalTo(2));
-        assertThat(params.get("index"), equalTo("index1"));
-        assertThat(params.get("docId"), equalTo("12"));
+        assertThat(params, equalTo(Map.of("index", "index1", "docId", "12")));
     }
 
     public void testEmptyPath() {
@@ -68,11 +62,11 @@ public class PathTrieTests extends ESTestCase {
 
         Map<String, String> params = new HashMap<>();
         assertThat(trie.retrieve("/a/test", params), equalTo("test1"));
-        assertThat(params.get("type"), equalTo("test"));
+        assertThat(params, equalTo(Map.of("type", "test")));
 
         params.clear();
         assertThat(trie.retrieve("/b/testX", params), equalTo("test2"));
-        assertThat(params.get("name"), equalTo("testX"));
+        assertThat(params, equalTo(Map.of("name", "testX")));
     }
 
     public void testSameNameOnDifferentPath() {
@@ -82,11 +76,11 @@ public class PathTrieTests extends ESTestCase {
 
         Map<String, String> params = new HashMap<>();
         assertThat(trie.retrieve("/a/c/test", params), equalTo("test1"));
-        assertThat(params.get("name"), equalTo("test"));
+        assertThat(params, equalTo(Map.of("name", "test")));
 
         params.clear();
         assertThat(trie.retrieve("/b/testX", params), equalTo("test2"));
-        assertThat(params.get("name"), equalTo("testX"));
+        assertThat(params, equalTo(Map.of("name", "testX")));
     }
 
     public void testPreferNonWildcardExecution() {
@@ -125,56 +119,31 @@ public class PathTrieTests extends ESTestCase {
         assertThat(trie.retrieve("/a", params, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED), equalTo("test1"));
         assertThat(trie.retrieve("/a", params, TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED), equalTo("test1"));
         assertThat(trie.retrieve("/a", params, TrieMatchingMode.WILDCARD_NODES_ALLOWED), equalTo("test1"));
-        Iterator<String> allPaths = trie.retrieveAll("/a", () -> params);
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo("test1"));
-        assertThat(allPaths.next(), equalTo("test1"));
-        assertThat(allPaths.next(), equalTo("test1"));
-        assertFalse(allPaths.hasNext());
+        assertThat(trie.retrieveAll("/a", () -> params).toList(), contains(null, "test1", "test1", "test1"));
 
         assertThat(trie.retrieve("/a/b", params, TrieMatchingMode.EXPLICIT_NODES_ONLY), nullValue());
         assertThat(trie.retrieve("/a/b", params, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED), equalTo("test4"));
         assertThat(trie.retrieve("/a/b", params, TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED), equalTo("test3"));
         assertThat(trie.retrieve("/a/b", params, TrieMatchingMode.WILDCARD_NODES_ALLOWED), equalTo("test3"));
-        allPaths = trie.retrieveAll("/a/b", () -> params);
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo("test4"));
-        assertThat(allPaths.next(), equalTo("test3"));
-        assertThat(allPaths.next(), equalTo("test3"));
-        assertFalse(allPaths.hasNext());
+        assertThat(trie.retrieveAll("/a/b", () -> params).toList(), contains(null, "test4", "test3", "test3"));
 
         assertThat(trie.retrieve("/a/b/c", params, TrieMatchingMode.EXPLICIT_NODES_ONLY), nullValue());
         assertThat(trie.retrieve("/a/b/c", params, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED), equalTo("test5"));
         assertThat(trie.retrieve("/a/b/c", params, TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED), equalTo("test7"));
         assertThat(trie.retrieve("/a/b/c", params, TrieMatchingMode.WILDCARD_NODES_ALLOWED), equalTo("test7"));
-        allPaths = trie.retrieveAll("/a/b/c", () -> params);
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo("test5"));
-        assertThat(allPaths.next(), equalTo("test7"));
-        assertThat(allPaths.next(), equalTo("test7"));
-        assertFalse(allPaths.hasNext());
+        assertThat(trie.retrieveAll("/a/b/c", () -> params).toList(), contains(null, "test5", "test7", "test7"));
 
         assertThat(trie.retrieve("/x/y/z", params, TrieMatchingMode.EXPLICIT_NODES_ONLY), nullValue());
         assertThat(trie.retrieve("/x/y/z", params, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED), nullValue());
         assertThat(trie.retrieve("/x/y/z", params, TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED), nullValue());
         assertThat(trie.retrieve("/x/y/z", params, TrieMatchingMode.WILDCARD_NODES_ALLOWED), equalTo("test9"));
-        allPaths = trie.retrieveAll("/x/y/z", () -> params);
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo("test9"));
-        assertFalse(allPaths.hasNext());
+        assertThat(trie.retrieveAll("/x/y/z", () -> params).toList(), contains(null, null, null, "test9"));
 
         assertThat(trie.retrieve("/d/e/f", params, TrieMatchingMode.EXPLICIT_NODES_ONLY), nullValue());
         assertThat(trie.retrieve("/d/e/f", params, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED), nullValue());
         assertThat(trie.retrieve("/d/e/f", params, TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED), nullValue());
         assertThat(trie.retrieve("/d/e/f", params, TrieMatchingMode.WILDCARD_NODES_ALLOWED), equalTo("test10"));
-        allPaths = trie.retrieveAll("/d/e/f", () -> params);
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo(null));
-        assertThat(allPaths.next(), equalTo("test10"));
-        assertFalse(allPaths.hasNext());
+        assertThat(trie.retrieveAll("/d/e/f", () -> params).toList(), contains(null, null, null, "test10"));
     }
 
     // https://github.com/elastic/elasticsearch/pull/17916
@@ -208,13 +177,11 @@ public class PathTrieTests extends ESTestCase {
 
         Map<String, String> params = new HashMap<>();
         assertThat(trie.retrieve("/a/b/c", params), equalTo("test1"));
-        assertThat(params.get("x"), equalTo("a"));
-        assertThat(params.get("y"), equalTo("b"));
-        assertThat(params.get("z"), equalTo("c"));
+        assertThat(params, equalTo(Map.of("x", "a", "y", "b", "z", "c")));
+
         params.clear();
         assertThat(trie.retrieve("/a/_y/c", params), equalTo("test2"));
-        assertThat(params.get("x"), equalTo("a"));
-        assertThat(params.get("k"), equalTo("c"));
+        assertThat(params, equalTo(Map.of("x", "a", "k", "c")));
     }
 
     public void testNamedWildcardAndLookupWithWildcard() {
@@ -227,23 +194,23 @@ public class PathTrieTests extends ESTestCase {
 
         Map<String, String> params = new HashMap<>();
         assertThat(trie.retrieve("/x/*", params), equalTo("test1"));
-        assertThat(params.get("test"), equalTo("*"));
+        assertThat(params, equalTo(Map.of("test", "*")));
 
-        params = new HashMap<>();
+        params.clear();
         assertThat(trie.retrieve("/b/a", params), equalTo("test2"));
-        assertThat(params.get("test"), equalTo("b"));
+        assertThat(params, equalTo(Map.of("test", "b")));
 
-        params = new HashMap<>();
+        params.clear();
         assertThat(trie.retrieve("/*", params), equalTo("test3"));
-        assertThat(params.get("test"), equalTo("*"));
+        assertThat(params, equalTo(Map.of("test", "*")));
 
-        params = new HashMap<>();
+        params.clear();
         assertThat(trie.retrieve("/*/_endpoint", params), equalTo("test4"));
-        assertThat(params.get("test"), equalTo("*"));
+        assertThat(params, equalTo(Map.of("test", "*")));
 
-        params = new HashMap<>();
+        params.clear();
         assertThat(trie.retrieve("a/*/_endpoint", params), equalTo("test5"));
-        assertThat(params.get("test"), equalTo("*"));
+        assertThat(params, equalTo(Map.of("test", "*")));
     }
 
     // https://github.com/elastic/elasticsearch/issues/14177
@@ -252,14 +219,12 @@ public class PathTrieTests extends ESTestCase {
         PathTrie<String> pathTrie = new PathTrie<>(RestUtils.REST_DECODER);
         pathTrie.insert("/{index}/{type}/{id}", "test");
         HashMap<String, String> params = new HashMap<>();
+
         assertThat(pathTrie.retrieve("/index/type/a%2Fe", params), equalTo("test"));
-        assertThat(params.get("index"), equalTo("index"));
-        assertThat(params.get("type"), equalTo("type"));
-        assertThat(params.get("id"), equalTo("a/e"));
+        assertThat(params, equalTo(Map.of("index", "index", "type", "type", "id", "a/e")));
+
         params.clear();
         assertThat(pathTrie.retrieve("/<logstash-{now%2Fd}>/type/id", params), equalTo("test"));
-        assertThat(params.get("index"), equalTo("<logstash-{now/d}>"));
-        assertThat(params.get("type"), equalTo("type"));
-        assertThat(params.get("id"), equalTo("id"));
+        assertThat(params, equalTo(Map.of("index", "<logstash-{now/d}>", "type", "type", "id", "id")));
     }
 }