1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.server.sdksandbox.verifier;
18 
19 import android.annotation.NonNull;
20 import android.annotation.Nullable;
21 
22 import java.util.Arrays;
23 import java.util.LinkedHashMap;
24 import java.util.List;
25 import java.util.Map;
26 
27 /**
28  * StringTrie is a string equality version of {@link com.android.tradefed.util.RegexTrie} It
29  * supports wildcard matching, null being the wildcard key. Wildcard can be in leaf nodes and inner
30  * nodes in the trie. This variant favors exact matching over wildcard matching. Matches by wildcard
31  * will be added to the captures. This structure is not thread safe.
32  *
33  * @param <V> type of the values contained in the trie
34  */
35 public class StringTrie<V> {
36     private V mValue = null;
37     private Map<String, StringTrie<V>> mChildren = new LinkedHashMap<String, StringTrie<V>>();
38 
39     /** Clears the trie */
clear()40     public void clear() {
41         mValue = null;
42         for (StringTrie<V> child : mChildren.values()) {
43             child.clear();
44         }
45         mChildren.clear();
46     }
47 
containsKey(String... strings)48     boolean containsKey(String... strings) {
49         return retrieve(strings) != null;
50     }
51 
52     @Nullable
recursivePut(V value, List<String> keys)53     V recursivePut(V value, List<String> keys) {
54         // Cases:
55         // 1) keys is empty -- set our value
56         // 2) keys is non-empty -- recurse downward, creating a child if necessary
57         if (keys.isEmpty()) {
58             V oldValue = mValue;
59             mValue = value;
60             return oldValue;
61         } else {
62             String curKey = keys.get(0);
63             List<String> nextKeys = keys.subList(1, keys.size());
64 
65             // Create a new child to handle
66             StringTrie<V> nextChild = mChildren.get(curKey);
67             if (nextChild == null) {
68                 nextChild = new StringTrie<V>();
69                 mChildren.put(curKey, nextChild);
70             }
71             return nextChild.recursivePut(value, nextKeys);
72         }
73     }
74 
75     /**
76      * Add an entry to the trie.
77      *
78      * @param value The value to set
79      * @param keys The sequence of {@link Strings}s that must be sequentially matched to retrieve
80      *     the associated {@code value}
81      * @return previous value stored for the given key sequence or null if no value was stored
82      */
put(V value, String... keys)83     public @Nullable V put(V value, String... keys) {
84         if (keys.length == 0) {
85             throw new IllegalArgumentException("string list must be non-empty.");
86         }
87         List<String> kList = Arrays.asList(keys);
88         return recursivePut(value, kList);
89     }
90 
91     @Nullable
recursiveRetrieve(List<String> captures, List<String> strings)92     V recursiveRetrieve(List<String> captures, List<String> strings) {
93         // Cases:
94         // 1) strings is empty -- return the value of this node
95         // 2) strings is non-empty -- find the first child that matches, recurse downward
96         if (strings.isEmpty()) {
97             return mValue;
98 
99         } else {
100             String curKey = strings.get(0);
101             List<String> nextKeys = strings.subList(1, strings.size());
102 
103             // a more specific match takes precedence over a wildcard match
104             if (mChildren.containsKey(curKey)) {
105                 V exactMatch = mChildren.get(curKey).recursiveRetrieve(captures, nextKeys);
106                 if (exactMatch != null) {
107                     return exactMatch;
108                 }
109             }
110 
111             V wildcardMatch = null;
112 
113             if (mChildren.containsKey(null)) {
114                 if (captures != null) {
115                     captures.add(curKey);
116                 }
117                 if (!nextKeys.isEmpty()) {
118                     // if there is a match after the wildcard continue down the tree
119                     if (mChildren.get(null).mChildren.containsKey(nextKeys.get(0))) {
120                         wildcardMatch = mChildren.get(null).recursiveRetrieve(captures, nextKeys);
121                     } else {
122                         // otherwise, keep matching wildcard
123                         wildcardMatch = recursiveRetrieve(captures, nextKeys);
124                     }
125                 } else {
126                     // last key to be matched with wildcard
127                     wildcardMatch = mChildren.get(null).recursiveRetrieve(captures, nextKeys);
128                 }
129             }
130 
131             return wildcardMatch;
132         }
133     }
134 
135     /**
136      * Fetch a value from the trie, by matching the provided sequence of {@link String}s to a
137      * sequence stored in the trie.
138      *
139      * @param strings A sequence of {@link String}s to match
140      * @return The associated value, or {@code null} if no value was found
141      */
retrieve(@onNull String... strings)142     public @Nullable V retrieve(@NonNull String... strings) {
143         return retrieve(null, strings);
144     }
145 
146     /**
147      * Fetch a value from the trie, by matching the provided sequence of {@link String}s to a
148      * sequence of {@link Strings}s stored in the trie. This version of the method also returns a
149      * {@link List} of matched captures.
150      *
151      * <p>For each level, the matched strings will be stored. Note that {@code captures} will be
152      * {@link List#clear()}ed before the retrieval begins. Also, if the retrieval fails after a
153      * partial sequence of matches, {@code captures} will still reflect the capture groups from the
154      * partial match.
155      *
156      * @param captures A {@code List<String>} through which wildcard matched keys will be returned.
157      * @param strings A sequence of {@link String}s to match
158      * @return The associated value, or {@code null} if no value was found
159      */
retrieve(@ullable List<String> captures, @NonNull String... strings)160     public @Nullable V retrieve(@Nullable List<String> captures, @NonNull String... strings) {
161         if (strings == null || strings.length == 0) {
162             throw new IllegalArgumentException("string list must be non-empty");
163         }
164         List<String> sList = Arrays.asList(strings);
165         if (captures != null) {
166             captures.clear();
167         }
168         return recursiveRetrieve(captures, sList);
169     }
170 
171     @Override
toString()172     public String toString() {
173         return String.format("{V: %s, C: %s}", mValue, mChildren);
174     }
175 }
176