Source code

001/*
002 * Copyright (c) 2016-2018 Daniel Ennis (Aikar) - MIT License
003 *
004 *  Permission is hereby granted, free of charge, to any person obtaining
005 *  a copy of this software and associated documentation files (the
006 *  "Software"), to deal in the Software without restriction, including
007 *  without limitation the rights to use, copy, modify, merge, publish,
008 *  distribute, sublicense, and/or sell copies of the Software, and to
009 *  permit persons to whom the Software is furnished to do so, subject to
010 *  the following conditions:
011 *
012 *  The above copyright notice and this permission notice shall be
013 *  included in all copies or substantial portions of the Software.
014 *
015 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
016 *  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
017 *  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
018 *  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
019 *  LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
020 *  OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
021 *  WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
022 */
023
024package co.aikar.util;
025
026import org.jetbrains.annotations.Nullable;
027
028import java.util.HashMap;
029import java.util.Iterator;
030import java.util.Map;
031import java.util.Objects;
032import java.util.Spliterators;
033import java.util.function.BiFunction;
034import java.util.function.Function;
035import java.util.function.Supplier;
036import java.util.stream.Stream;
037import java.util.stream.StreamSupport;
038
039public class Table <R, C, V> implements Iterable<Table.Entry<R, C, V>>{
040
041    private final Map<R, Map<C, V>> rowMap;
042    private final Function<R, Map<C, V>> colMapSupplier;
043
044    public Table() {
045        this(new HashMap<>(), (Supplier<Map<C, V>>) HashMap::new);
046    }
047
048    public Table(Supplier<Map<C, V>> columnMapSupplier) {
049        this(new HashMap<>(), columnMapSupplier);
050    }
051
052    public Table(Map<R, Map<C, V>> backingRowMap, Supplier<Map<C, V>> columnMapSupplier) {
053        this(backingRowMap, (r) -> columnMapSupplier.get());
054    }
055
056    public Table(Map<R, Map<C, V>> backingRowMap, Function<R, Map<C, V>> columnMapSupplier) {
057        this.rowMap = backingRowMap;
058        this.colMapSupplier = columnMapSupplier;
059    }
060
061
062    public V get(R row, C col) {
063        return getIfExists(row, col);
064    }
065
066    public V getOrDefault(R row, C col, V def) {
067        Map<C, V> colMap = getColMapIfExists(row);
068        if (colMap == null) {
069            return def;
070        }
071
072        V v = colMap.get(col);
073        if (v != null || colMap.containsKey(col)) {
074            return v;
075        }
076        return def;
077    }
078
079    public boolean containsKey(R row, C col) {
080        Map<C, V> colMap = getColMapIfExists(row);
081        if (colMap == null) {
082            return false;
083        }
084        return colMap.containsKey(col);
085    }
086
087    @Nullable
088    public V put(R row, C col, V val) {
089        return getColMapForWrite(row).put(col, val);
090    }
091
092    public void forEach(TableConsumer<R, C, V> consumer) {
093        for (Iterator<Entry<R, C, V>> it = this.iterator(); it.hasNext(); ) {
094            Entry<R, C, V> entry = it.next();
095            consumer.accept(entry.getRow(), entry.getCol(), entry.getValue());
096        }
097    }
098
099    public void forEach(TablePredicate<R, C, V> predicate) {
100        for (Iterator<Entry<R, C, V>> it = this.iterator(); it.hasNext(); ) {
101            Entry<R, C, V> entry = it.next();
102            if (!predicate.test(entry.getRow(), entry.getCol(), entry.getValue())) {
103                return;
104            }
105        }
106    }
107
108    public void removeIf(TablePredicate<R, C, V> predicate) {
109        for (Iterator<Entry<R, C, V>> it = this.iterator(); it.hasNext(); ) {
110            Entry<R, C, V> entry = it.next();
111            if (predicate.test(entry.getRow(), entry.getCol(), entry.getValue())) {
112                it.remove();
113            }
114        }
115    }
116
117    public Stream<Entry<R, C, V>> stream() {
118        return stream(false);
119    }
120
121    public Stream<Entry<R, C, V>> stream(boolean parallel) {
122        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator(), 0), parallel);
123    }
124
125    public Iterator<Entry<R, C, V>> iterator() {
126        return new Iterator<Entry<R, C, V>>() {
127            Iterator<Map.Entry<R, Map<C, V>>> rowIter = rowMap.entrySet().iterator();
128            Iterator<Map.Entry<C, V>> colIter = null;
129
130            private Map.Entry<R, Map<C, V>> rowEntry;
131            private Map.Entry<C, V> colEntry;
132
133            private Entry<R, C, V> next = getNext();
134
135            private Entry<R, C, V> getNext() {
136                if (colIter == null || !colIter.hasNext()) {
137                    if (!rowIter.hasNext()) {
138                        return null;
139                    }
140                    rowEntry = rowIter.next();
141                    colIter = rowEntry.getValue().entrySet().iterator();
142                }
143
144                if (!colIter.hasNext()) {
145                    return null;
146                }
147
148                colEntry = colIter.next();
149
150                return new Node(rowEntry, colEntry);
151            }
152
153            @Override
154            public boolean hasNext() {
155                return this.next != null;
156            }
157
158            @Override
159            public Entry<R, C, V> next() {
160                Entry<R, C, V> entry = this.next;
161                this.next = getNext();
162                return entry;
163            }
164
165            @Override
166            public void remove() {
167                this.colIter.remove();
168            }
169        };
170    }
171
172    public void replaceAll(TableFunction<R, C, V, V> function) {
173        for (Iterator<Entry<R, C, V>> it = this.iterator(); it.hasNext(); ) {
174            Entry<R, C, V> entry = it.next();
175            entry.setValue(function.compose(entry.getRow(), entry.getCol(), entry.getValue()));
176        }
177    }
178
179    public V remove(R row, C col) {
180        Map<C, V> rowMap = this.rowMap.get(row);
181        if (rowMap == null) {
182            return null;
183        }
184        return rowMap.remove(col);
185    }
186
187    @Nullable
188    public V replace(R row, C col, V val) {
189        Map<C, V> rowMap = getColMapIfExists(row);
190        if (rowMap == null) {
191            return null;
192        }
193        if (rowMap.get(col) != null || rowMap.containsKey(col)) {
194            return rowMap.put(col, val);
195        }
196        return null;
197    }
198
199
200    @Nullable
201    public boolean replace(R row, C col, V old, V val) {
202        Map<C, V> rowMap = getColMapIfExists(row);
203        if (rowMap == null) {
204            return false;
205        }
206        if (Objects.equals(rowMap.get(col), old)) {
207            rowMap.put(col, val);
208            return true;
209        }
210        return false;
211    }
212
213    public V computeIfAbsent(R row, C col, BiFunction<R, C, V> function) {
214        return getColMapForWrite(row).computeIfAbsent(col, c -> function.apply(row, col));
215    }
216
217    public V computeIfPresent(R row, C col, TableFunction<R, C, V, V> function) {
218        Map<C, V> colMap = getColMapForWrite(row);
219        V v = colMap.computeIfPresent(col, (c, old) -> function.compose(row, col, old));
220        removeIfEmpty(row, colMap);
221        return v;
222    }
223
224    public V compute(R row, C col, TableFunction<R, C, V, V> function) {
225        Map<C, V> colMap = getColMapForWrite(row);
226        V v = colMap.compute(col, (c, old) -> function.compose(row, col, old));
227        removeIfEmpty(row, colMap);
228        return v;
229    }
230
231    public V merge(R row, C col, V val, TableFunction<R, C, V, V> function) {
232        Map<C, V> colMap = getColMapForWrite(row);
233        V v = colMap.merge(col, val, (c, old) -> function.compose(row, col, old));
234        removeIfEmpty(row, colMap);
235        return v;
236    }
237
238    public Map<C, V> row(R row) {
239        Map<C, V> EMPTY = new HashMap<>(0);
240        return new DelegatingMap<C, V>() {
241            @Override
242            public Map<C, V> delegate(boolean readOnly) {
243                if (readOnly) {
244                    return Table.this.rowMap.getOrDefault(row, EMPTY);
245                }
246                return getColMapForWrite(row);
247            }
248
249            @Override
250            public V remove(Object key) {
251                Map<C, V> delegate = delegate(false);
252                V remove = delegate.remove(key);
253                removeIfEmpty(row, delegate);
254                return remove;
255            }
256            // iterators may leave us empty, but the next get will remove it.
257        };
258    }
259    // Other stuff
260
261    public interface TablePredicate<R, C, V> {
262        boolean test(R row, C col, V val);
263    }
264    public interface TableFunction<R, C, V, RETURN> {
265        RETURN compose(R row, C col, V val);
266    }
267    public interface TableConsumer<R, C, V> {
268        void accept(R row, C col, V val);
269    }
270
271    private V getIfExists(R row, C col) {
272        Map<C, V> colMap = getColMapIfExists(row);
273        if (colMap == null) {
274            return null;
275        }
276
277        return colMap.get(col);
278    }
279
280    private Map<C, V> getColMapIfExists(R row) {
281        Map<C, V> colMap = this.rowMap.get(row);
282        if (colMap != null && colMap.isEmpty()) {
283            rowMap.remove(row);
284            colMap = null;
285        }
286        return colMap;
287    }
288
289    private Map<C, V> getColMapForWrite(R row) {
290        return this.rowMap.computeIfAbsent(row, this.colMapSupplier);
291    }
292
293    private void removeIfEmpty(R row, Map<C, V> colMap) {
294        if (colMap.isEmpty()) {
295            this.rowMap.remove(row);
296        }
297    }
298
299    public interface Entry <R, C, V> {
300        R getRow();
301        C getCol();
302        V getValue();
303        V setValue(V value);
304    }
305
306    private class Node implements Entry <R, C, V> {
307
308        private final Map.Entry<R, Map<C, V>> rowEntry;
309        private final Map.Entry<C, V> colEntry;
310
311        Node(Map.Entry<R, Map<C, V>> rowEntry, Map.Entry<C, V> entry) {
312            this.rowEntry = rowEntry;
313            this.colEntry = entry;
314        }
315
316        @Override
317        public final R getRow() {
318            return rowEntry.getKey();
319        }
320
321        @Override
322        public final C getCol() {
323            return colEntry.getKey();
324        }
325
326        @Override
327        public final V getValue() {
328            return colEntry.getValue();
329        }
330
331        @Override
332        public final V setValue(V value) {
333            return colEntry.setValue(value);
334        }
335    }
336}