1#
2# set.py
3#
4# This source file is part of the FoundationDB open source project
5#
6# Copyright 2013-2018 Apple Inc. and the FoundationDB project authors
7#
8# Licensed under the Apache License, Version 2.0 (the "License");
9# you may not use this file except in compliance with the License.
10# You may obtain a copy of the License at
11#
12#     http://www.apache.org/licenses/LICENSE-2.0
13#
14# Unless required by applicable law or agreed to in writing, software
15# distributed under the License is distributed on an "AS IS" BASIS,
16# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17# See the License for the specific language governing permissions and
18# limitations under the License.
19#
20
21import fdb
22import fdb.tuple
23
24fdb.api_version(16)
25
26
27def nextStopToNone(gen):
28    try:
29        return gen.next()
30    except StopIteration:
31        return None
32
33
34class FdbSet (object):
35    def __init__(self, path):
36        self._path = path
37
38    @fdb.transactional
39    def length(self, tr):
40        setLength = 0
41        for k, v in tr[fdb.tuple.range((self._path,))]:
42            setLength += 1
43        return setLength
44
45    @fdb.transactional
46    def iterate(self, tr):
47        for k, v in tr[fdb.tuple.range((self._path,))]:
48            yield fdb.tuple.unpack(k)[1]
49
50    @fdb.transactional
51    def contains(self, tr, x):
52        return tr[fdb.tuple.pack((self._path, x))].present()
53
54    @fdb.transactional
55    def issubset(self, tr, t):  # s <= t
56        for k, v in tr[fdb.tuple.range((self._path,))]:
57            if not t.contains(tr, fdb.tuple.unpack(k)[1]):
58                return False
59        return True
60
61    @fdb.transactional
62    def issuperset(self, tr, t):  # s >= t
63        return t.issubset(tr, self)
64
65    @fdb.transactional
66    def union(self, tr, t):  # s | t
67        s_gen = self.iterate(tr)
68        t_gen = t.iterate(tr)
69        s_key = nextStopToNone(s_gen)
70        t_key = nextStopToNone(t_gen)
71
72        while True:
73            if t_key == None and s_key == None:
74                return
75            elif t_key == None or (s_key != None and s_key < t_key):
76                yield s_key
77                s_key = nextStopToNone(s_gen)
78            elif s_key == None or s_key > t_key:
79                yield t_key
80                t_key = nextStopToNone(t_gen)
81            else:
82                yield s_key
83                s_key = nextStopToNone(s_gen)
84                t_key = nextStopToNone(t_gen)
85
86    @fdb.transactional
87    def intersection(self, tr, t):  # s & t
88        s_key = self.first_greater_or_equal(tr, "")
89        t_key = t.first_greater_or_equal(tr, "")
90
91        while True:
92            if t_key == None or s_key == None:
93                return
94            elif s_key < t_key:
95                s_key = self.first_greater_or_equal(tr, t_key)
96            elif s_key > t_key:
97                t_key = t.first_greater_or_equal(tr, s_key)
98            else:
99                yield s_key
100                s_key = self.first_greater_than(tr, s_key)
101                t_key = t.first_greater_than(tr, t_key)
102
103    @fdb.transactional
104    def difference(self, tr, t):  # s - t
105        s_gen = self.iterate(tr)
106        s_key = nextStopToNone(s_gen)
107        t_key = t.first_greater_or_equal(tr, "")
108
109        while True:
110            if s_key == None:
111                return
112            elif t_key == None or s_key < t_key:
113                yield s_key
114                s_key = nextStopToNone(s_gen)
115            elif s_key > t_key:
116                t_key = t.first_greater_or_equal(tr, s_key)
117            else:
118                s_key = nextStopToNone(s_gen)
119
120    @fdb.transactional
121    def symmetric_difference(self, tr, t):  # s ^ t
122        s_gen = self.iterate(tr)
123        t_gen = t.iterate(tr)
124        s_key = nextStopToNone(s_gen)
125        t_key = nextStopToNone(t_gen)
126
127        while True:
128            if t_key == None and s_key == None:
129                return
130            elif t_key == None or (s_key != None and s_key < t_key):
131                yield s_key
132                s_key = nextStopToNone(s_gen)
133            elif s_key == None or s_key > t_key:
134                yield t_key
135                t_key = nextStopToNone(t_gen)
136            else:
137                s_key = nextStopToNone(s_gen)
138                t_key = nextStopToNone(t_gen)
139
140    @fdb.transactional
141    def update(self, tr, t):  # s |= t T
142        for k in t.iterate(tr):
143            self.add(tr, k)
144
145    @fdb.transactional
146    def intersection_update(self, tr, t):  # s &= t
147        lastValue = fdb.tuple.pack((self._path,))
148        for k in self.intersection(tr, t):
149            if k != lastValue:
150                del tr[lastValue + '\x00':fdb.tuple.pack((self._path, k))]
151            lastValue = fdb.tuple.pack((self._path, k))
152        del tr[lastValue + '\x00':fdb.tuple.pack((self._path + chr(0),))]
153
154    @fdb.transactional
155    def difference_update(self, tr, t):  # s -= t
156        for k in self.intersection(tr, t):
157            del tr[fdb.tuple.pack((self._path, k))]
158
159    @fdb.transactional
160    def symmetric_difference_update(self, tr, t):  # s ^ t
161        s_gen = self.iterate(tr)
162        t_gen = t.iterate(tr)
163        s_key = nextStopToNone(s_gen)
164        t_key = nextStopToNone(t_gen)
165
166        while True:
167            if t_key == None and s_key == None:
168                return
169            elif t_key == None or (s_key != None and s_key < t_key):
170                s_key = nextStopToNone(s_gen)
171            elif s_key == None or s_key > t_key:
172                self.add(tr, t_key)
173                t_key = nextStopToNone(t_gen)
174            else:
175                remove_key = s_key
176                s_key = nextStopToNone(s_gen)
177                t_key = nextStopToNone(t_gen)
178                self.remove(tr, remove_key)
179
180    @fdb.transactional
181    def add(self, tr, x):
182        tr[fdb.tuple.pack((self._path, x))] = ""
183
184    @fdb.transactional
185    def remove(self, tr, x):
186        if tr[fdb.tuple.pack((self._path, x))] == None:
187            raise KeyError
188        del tr[fdb.tuple.pack((self._path, x))]
189
190    @fdb.transactional
191    def discard(self, tr, x):
192        del tr[fdb.tuple.pack((self._path, x))]
193
194    @fdb.transactional
195    def pop(self, tr):
196        key = tr.get_key(fdb.KeySelector.first_greater_or_equal(self._path))
197        if self._keyInRange(key):
198            del tr[key]
199            return fdb.tuple.unpack(key)[1]
200        raise KeyError
201
202    @fdb.transactional
203    def clear(self, tr):
204        tr.clear_range_startswith(self._path)
205
206    @fdb.transactional
207    def first_greater_than(self, tr, x):
208        key = tr.get_key(fdb.KeySelector.first_greater_than(fdb.tuple.pack((self._path, x))))
209        if self._keyInRange(key):
210            return fdb.tuple.unpack(key)[1]
211        return None
212
213    @fdb.transactional
214    def first_greater_or_equal(self, tr, x):
215        key = tr.get_key(fdb.KeySelector.first_greater_or_equal(fdb.tuple.pack((self._path, x))))
216        if self._keyInRange(key):
217            return fdb.tuple.unpack(key)[1]
218        return None
219
220    def _keyInRange(self, key):
221        return key < fdb.tuple.pack((self._path + chr(0),))
222
223
224def test(db):
225    print "starting set test"
226    tr = db.create_transaction()
227
228    del tr[:]
229
230    a = FdbSet("a")
231    a.add(tr, "apple")
232    a.add(tr, "banana")
233    a.add(tr, "orange")
234
235    b = FdbSet("b")
236    b.add(tr, "banana")
237    b.add(tr, "grape")
238    b.add(tr, "strawberry")
239
240    c = FdbSet("c")
241    c.add(tr, "grape")
242
243    print "set a:"
244    for k in a.iterate(tr):
245        print k
246
247    print "set b:"
248    for k in b.iterate(tr):
249        print k
250
251    print "set c:"
252    for k in c.iterate(tr):
253        print k
254
255    print "b contains strawberry: {0}".format(b.contains(tr, "strawberry"))
256    print "remove strawberry from b"
257
258    b.remove(tr, "strawberry")
259
260    print "b contains strawberry: {0}".format(b.contains(tr, "strawberry"))
261
262    try:
263        b.remove(tr, "strawberry")
264        print "survived second remove of strawberry"
265    except KeyError:
266        print "failed second remove of strawberry"
267
268    print "insert strawberry into b"
269
270    b.add(tr, "strawberry")
271
272    print "b contains strawberry: {0}".format(b.contains(tr, "strawberry"))
273
274    print "discard strawberry from b"
275
276    b.discard(tr, "strawberry")
277
278    print "b contains strawberry: {0}".format(b.contains(tr, "strawberry"))
279
280    b.discard(tr, "strawberry")
281
282    print "b length: {0}".format(b.length(tr))
283
284    print "c issubset a: {0}".format(c.issubset(tr, a))
285
286    print "c issubset b: {0}".format(c.issubset(tr, b))
287
288    print "a union b:"
289    for k in a.union(tr, b):
290        print k
291
292    print "a intersection b:"
293    for k in a.intersection(tr, b):
294        print k
295
296    print "a difference b:"
297    for k in a.difference(tr, b):
298        print k
299
300    print "b difference a:"
301    for k in b.difference(tr, a):
302        print k
303
304    print "a symmetric_difference b:"
305    for k in a.symmetric_difference(tr, b):
306        print k
307
308    print "a update c:"
309    a.update(tr, c)
310    for k in a.iterate(tr):
311        print k
312
313    print "b intersection_update a:"
314    b.intersection_update(tr, a)
315    for k in b.iterate(tr):
316        print k
317
318    print "a difference_update c"
319    a.difference_update(tr, c)
320    for k in a.iterate(tr):
321        print k
322
323    print "b symmetric_difference_update a"
324    b.symmetric_difference_update(tr, a)
325    for k in b.iterate(tr):
326        print k
327
328    print "popping items from a"
329    try:
330        while True:
331            print "popped {0}".format(a.pop(tr))
332    except KeyError:
333        print "finished popping"
334
335    print "a length: {0}".format(a.length(tr))
336
337    print "clearing b"
338
339    b.clear(tr)
340
341    print "b length: {0}".format(b.length(tr))
342
343
344db = fdb.open()
345test(db)
346