s2.util.container.rbtree

1 module s2.util.container.rbtree_map;
2 
3 import std.container.rbtree;
4 import std.algorithm : map;
5 import std.functional : binaryFun;
6 import std.meta : allSatisfy;
7 import std.range : ElementType, isInputRange;
8 import std.traits : isDynamicArray, isImplicitlyConvertible;
9 
10 /**
11  * A dictionary or associative array backed by a Red-Black tree.
12  */
13 final class RBTreeMap(KeyT, ValueT, alias KeyLessF = "a < b", bool allowDuplicates = false) {
14 public:
15   static struct Pair {
16     KeyT key;
17     ValueT value;
18   }
19 
20   alias keyLess = binaryFun!KeyLessF;
21 
22   alias RedBlackTreeT =
23       RedBlackTree!(Pair, (pair1, pair2) => keyLess(pair1.key, pair2.key), allowDuplicates);
24 
25   RedBlackTreeT rbTree;
26 
27   // Forward compatible methods like: empty(), length(), opSlice(), etc.
28   alias rbTree this;
29 
30   this() {
31     rbTree = new RedBlackTreeT();
32   }
33 
34   this(Pair[] elems...) {
35     rbTree = new RedBlackTreeT(elems);
36   }
37 
38   this(PairRange)(PairRange pairRange)
39   if (isInputRange!PairRange && isImplicitlyConvertible!(ElementType!PairRange, Pair)) {
40     rbTree = new RedBlackTreeT(pairRange);
41   }
42 
43   override
44   bool opEquals(Object rhs) {
45     RBTreeMap that = cast(RBTreeMap) rhs;
46     if (that is null) return false;
47 
48     return rbTree == that.rbTree;
49   }
50 
51   /// Insertion
52   size_t stableInsert(K, V)(K key, V value)
53   if (isImplicitlyConvertible!(K, KeyT) && isImplicitlyConvertible!(V, ValueT)) {
54     return rbTree.stableInsert(Pair(key, value));
55   }
56   alias insert = stableInsert;
57 
58   ValueT opIndexAssign(ValueT value, KeyT key) {
59     rbTree.stableInsert(Pair(key, value));
60     return value;
61   }
62 
63   /// Membership
64   bool opBinaryRight(string op)(KeyT key) const
65   if (op == "in") {
66     return Pair(key) in rbTree;
67   }
68 
69   /// Removal
70   size_t removeKey(K...)(K keys)
71   if (allSatisfy!(isImplicitlyConvertibleToKey, K)) {
72     KeyT[K.length] toRemove = [keys];
73     return removeKey(toRemove[]);
74   }
75 
76   //Helper for removeKey.
77   private template isImplicitlyConvertibleToKey(K)
78   {
79     enum isImplicitlyConvertibleToKey = isImplicitlyConvertible!(K, KeyT);
80   }
81 
82   size_t removeKey(K)(K[] keys)
83   if (isImplicitlyConvertible!(K, KeyT)) {
84     auto keyPairs = keys.map!(key => Pair(key));
85     return rbTree.removeKey(keyPairs);
86   }
87 
88   size_t removeKey(KeyRange)(KeyRange keyRange)
89   if (isInputRange!KeyRange
90       && isImplicitlyConvertible!(ElementType!KeyRange, KeyT)
91       && !isDynamicArray!KeyRange) {
92     auto keyPairs = keys.map(key => Pair(key));
93     return rbTree.removeKey(keyPairs);
94   }
95 
96   /// Ranges
97   RedBlackTreeT.Range upperBound(KeyT key) {
98     return rbTree.upperBound(Pair(key));
99   }
100 
101   RedBlackTreeT.ConstRange upperBound(KeyT key) const {
102     return rbTree.upperBound(Pair(key));
103   }
104 
105   RedBlackTreeT.ImmutableRange upperBound(KeyT key) immutable {
106     return rbTree.upperBound(Pair(key));
107   }
108 
109   RedBlackTreeT.Range lowerBound(KeyT key) {
110     return rbTree.lowerBound(Pair(key));
111   }
112 
113   RedBlackTreeT.ConstRange lowerBound(KeyT key) const {
114     return rbTree.lowerBound(Pair(key));
115   }
116 
117   RedBlackTreeT.ImmutableRange lowerBound(KeyT key) immutable {
118     return rbTree.lowerBound(Pair(key));
119   }
120 
121   auto equalRange(KeyT key) {
122     return rbTree.equalRange(Pair(key));
123   }
124 
125 }