1 /*
2  * Copyright (C) 2016 Square, Inc.
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 package okio;
17
18 import java.util.AbstractList;
19 import java.util.ArrayList;
20 import java.util.Arrays;
21 import java.util.Collections;
22 import java.util.List;
23 import java.util.RandomAccess;
24
25 /** An indexed set of values that may be read with {@link BufferedSource#select}. */
26 public final class Options extends AbstractList<ByteString> implements RandomAccess {
27   final ByteString[] byteStrings;
28   final int[] trie;
29
30   private Options(ByteString[] byteStrings, int[] trie) {
31     this.byteStrings = byteStrings;
32     this.trie = trie;
33   }
34
35   public static Options of(ByteString... byteStrings) {
36     if (byteStrings.length == 0) {
37       // With no choices we must always return -1. Create a trie that selects from an empty set.
38       return new Options(new ByteString[0], new int[] { 0, -1 });
39     }
40
41     // Sort the byte strings which is required when recursively building the trie. Map the sorted
42     // indexes to the caller's indexes.
43     List<ByteString> list = new ArrayList<>(Arrays.asList(byteStrings));
44     Collections.sort(list);
45     List<Integer> indexes = new ArrayList<>();
46     for (int i = 0; i < list.size(); i++) {
47       indexes.add(-1);
48     }
49     for (int i = 0; i < list.size(); i++) {
50       int sortedIndex = Collections.binarySearch(list, byteStrings[i]);
51       indexes.set(sortedIndex, i);
52     }
53     if (list.get(0).size() == 0) {
54       throw new IllegalArgumentException("the empty byte string is not a supported option");
55     }
56
57     // Strip elements that will never be returned because they follow their own prefixes. For
58     // example, if the caller provides ["abc", "abcde"] we will never return "abcde" because we
59     // return as soon as we encounter "abc".
60     for (int a = 0; a < list.size(); a++) {
61       ByteString prefix = list.get(a);
62       for (int b = a + 1; b < list.size(); ) {
63         ByteString byteString = list.get(b);
64         if (!byteString.startsWith(prefix)) break;
65         if (byteString.size() == prefix.size()) {
66           throw new IllegalArgumentException("duplicate option: " + byteString);
67         }
68         if (indexes.get(b) > indexes.get(a)) {
69           list.remove(b);
70           indexes.remove(b);
71         } else {
72           b++;
73         }
74       }
75     }
76
77     Buffer trieBytes = new Buffer();
78     buildTrieRecursive(0L, trieBytes, 0, list, 0, list.size(), indexes);
79
80     int[] trie = new int[intCount(trieBytes)];
81     for (int i = 0; i < trie.length; i++) {
82       trie[i] = trieBytes.readInt();
83     }
84     if (!trieBytes.exhausted()) {
85       throw new AssertionError();
86     }
87
88     return new Options(byteStrings.clone() /* Defensive copy. */, trie);
89   }
90
91   /**
92    * Builds a trie encoded as an int array. Nodes in the trie are of two types: SELECT and SCAN.
93    *
94    * SELECT nodes are encoded as:
95    *  - selectChoiceCount: the number of bytes to choose between (a positive int)
96    *  - prefixIndex: the result index at the current position or -1 if the current position is not
97    *    a result on its own
98    *  - a sorted list of selectChoiceCount bytes to match against the input string
99    *  - a heterogeneous list of selectChoiceCount result indexes (>= 0) or offsets (< 0) of the
100    *    next node to follow. Elements in this list correspond to elements in the preceding list.
101    *    Offsets are negative and must be multiplied by -1 before being used.
102    *
103    * SCAN nodes are encoded as:
104    *  - scanByteCount: the number of bytes to match in sequence. This count is negative and must
105    *    be multiplied by -1 before being used.
106    *  - prefixIndex: the result index at the current position or -1 if the current position is not
107    *    a result on its own
108    *  - a list of scanByteCount bytes to match
109    *  - nextStep: the result index (>= 0) or offset (< 0) of the next node to follow. Offsets are
110    *    negative and must be multiplied by -1 before being used.
111    *
112    * This structure is used to improve locality and performance when selecting from a list of
113    * options.
114    */
115   private static void buildTrieRecursive(
116       long nodeOffset,
117       Buffer node,
118       int byteStringOffset,
119       List<ByteString> byteStrings,
120       int fromIndex,
121       int toIndex,
122       List<Integer> indexes) {
123     if (fromIndex >= toIndex) throw new AssertionError();
124     for (int i = fromIndex; i < toIndex; i++) {
125       if (byteStrings.get(i).size() < byteStringOffset) throw new AssertionError();
126     }
127
128     ByteString from = byteStrings.get(fromIndex);
129     ByteString to = byteStrings.get(toIndex - 1);
130     int prefixIndex = -1;
131
132     // If the first element is already matched, that's our prefix.
133     if (byteStringOffset == from.size()) {
134       prefixIndex = indexes.get(fromIndex);
135       fromIndex++;
136       from = byteStrings.get(fromIndex);
137     }
138
139     if (from.getByte(byteStringOffset) != to.getByte(byteStringOffset)) {
140       // If we have multiple bytes to choose from, encode a SELECT node.
141       int selectChoiceCount = 1;
142       for (int i = fromIndex + 1; i < toIndex; i++) {
143         if (byteStrings.get(i - 1).getByte(byteStringOffset)
144             != byteStrings.get(i).getByte(byteStringOffset)) {
145           selectChoiceCount++;
146         }
147       }
148
149       // Compute the offset that childNodes will get when we append it to node.
150       long childNodesOffset = nodeOffset + intCount(node) + 2 + (selectChoiceCount * 2);
151
152       node.writeInt(selectChoiceCount);
153       node.writeInt(prefixIndex);
154
155       for (int i = fromIndex; i < toIndex; i++) {
156         byte rangeByte = byteStrings.get(i).getByte(byteStringOffset);
157         if (i == fromIndex || rangeByte != byteStrings.get(i - 1).getByte(byteStringOffset)) {
158           node.writeInt(rangeByte & 0xff);
159         }
160       }
161
162       Buffer childNodes = new Buffer();
163       int rangeStart = fromIndex;
164       while (rangeStart < toIndex) {
165         byte rangeByte = byteStrings.get(rangeStart).getByte(byteStringOffset);
166         int rangeEnd = toIndex;
167         for (int i = rangeStart + 1; i < toIndex; i++) {
168           if (rangeByte != byteStrings.get(i).getByte(byteStringOffset)) {
169             rangeEnd = i;
170             break;
171           }
172         }
173
174         if (rangeStart + 1 == rangeEnd
175             && byteStringOffset + 1 == byteStrings.get(rangeStart).size()) {
176           // The result is a single index.
177           node.writeInt(indexes.get(rangeStart));
178         } else {
179           // The result is another node.
180           node.writeInt((int) (-1 * (childNodesOffset + intCount(childNodes))));
181           buildTrieRecursive(
182               childNodesOffset,
183               childNodes,
184               byteStringOffset + 1,
185               byteStrings,
186               rangeStart,
187               rangeEnd,
188               indexes);
189         }
190
191         rangeStart = rangeEnd;
192       }
193
194       node.write(childNodes, childNodes.size());
195
196     } else {
197       // If all of the bytes are the same, encode a SCAN node.
198       int scanByteCount = 0;
199       for (int i = byteStringOffset, max = Math.min(from.size(), to.size()); i < max; i++) {
200         if (from.getByte(i) == to.getByte(i)) {
201           scanByteCount++;
202         } else {
203           break;
204         }
205       }
206
207       // Compute the offset that childNodes will get when we append it to node.
208       long childNodesOffset = nodeOffset + intCount(node) + 2 + scanByteCount + 1;
209
210       node.writeInt(-scanByteCount);
211       node.writeInt(prefixIndex);
212
213       for (int i = byteStringOffset; i < byteStringOffset + scanByteCount; i++) {
214         node.writeInt(from.getByte(i) & 0xff);
215       }
216
217       if (fromIndex + 1 == toIndex) {
218         // The result is a single index.
219         if (byteStringOffset + scanByteCount != byteStrings.get(fromIndex).size()) {
220           throw new AssertionError();
221         }
222         node.writeInt(indexes.get(fromIndex));
223       } else {
224         // The result is another node.
225         Buffer childNodes = new Buffer();
226         node.writeInt((int) (-1 * (childNodesOffset + intCount(childNodes))));
227         buildTrieRecursive(
228             childNodesOffset,
229             childNodes,
230             byteStringOffset + scanByteCount,
231             byteStrings,
232             fromIndex,
233             toIndex,
234             indexes);
235         node.write(childNodes, childNodes.size());
236       }
237     }
238   }
239
240   @Override public ByteString get(int i) {
241     return byteStrings[i];
242   }
243
244   @Override public final int size() {
245     return byteStrings.length;
246   }
247
248   private static int intCount(Buffer trieBytes) {
249     return (int) (trieBytes.size() / 4);
250   }
251 }
252