1//
2// immer: immutable data structures for C++
3// Copyright (C) 2016, 2017, 2018 Juan Pedro Bolivar Puente
4//
5// This software is distributed under the Boost Software License, Version 1.0.
6// See accompanying file LICENSE or copy at http://boost.org/LICENSE_1_0.txt
7//
8
9#pragma once
10
11#include <immer/detail/rbts/node.hpp>
12#include <immer/detail/rbts/position.hpp>
13#include <immer/detail/rbts/operations.hpp>
14
15#include <immer/detail/type_traits.hpp>
16
17#include <cassert>
18#include <memory>
19#include <numeric>
20
21namespace immer {
22namespace detail {
23namespace rbts {
24
25template <typename T,
26 typename MemoryPolicy,
27 bits_t B,
28 bits_t BL>
29struct rbtree
30{
31 using node_t = node<T, MemoryPolicy, B, BL>;
32 using edit_t = typename node_t::edit_t;
33 using owner_t = typename MemoryPolicy::transience_t::owner;
34
35 size_t size;
36 shift_t shift;
37 node_t* root;
38 node_t* tail;
39
40 static const rbtree& empty()
41 {
42 static const rbtree empty_ {
43 0,
44 BL,
45 node_t::make_inner_n(0u),
46 node_t::make_leaf_n(0u)
47 };
48 return empty_;
49 }
50
51 template <typename U>
52 static auto from_initializer_list(std::initializer_list<U> values)
53 {
54 auto e = owner_t{};
55 auto result = rbtree{empty()};
56 for (auto&& v : values)
57 result.push_back_mut(e, v);
58 return result;
59 }
60
61 template <typename Iter, typename Sent,
62 std::enable_if_t
63 <compatible_sentinel_v<Iter, Sent>, bool> = true>
64 static auto from_range(Iter first, Sent last)
65 {
66 auto e = owner_t{};
67 auto result = rbtree{empty()};
68 for (; first != last; ++first)
69 result.push_back_mut(e, *first);
70 return result;
71 }
72
73 static auto from_fill(size_t n, T v)
74 {
75 auto e = owner_t{};
76 auto result = rbtree{empty()};
77 while (n --> 0)
78 result.push_back_mut(e, v);
79 return result;
80 }
81
82 rbtree(size_t sz, shift_t sh, node_t* r, node_t* t)
83 : size{sz}, shift{sh}, root{r}, tail{t}
84 {
85 assert(check_tree());
86 }
87
88 rbtree(const rbtree& other)
89 : rbtree{other.size, other.shift, other.root, other.tail}
90 {
91 inc();
92 }
93
94 rbtree(rbtree&& other)
95 : rbtree{empty()}
96 {
97 swap(*this, other);
98 }
99
100 rbtree& operator=(const rbtree& other)
101 {
102 auto next = other;
103 swap(*this, next);
104 return *this;
105 }
106
107 rbtree& operator=(rbtree&& other)
108 {
109 swap(*this, other);
110 return *this;
111 }
112
113 friend void swap(rbtree& x, rbtree& y)
114 {
115 using std::swap;
116 swap(x.size, y.size);
117 swap(x.shift, y.shift);
118 swap(x.root, y.root);
119 swap(x.tail, y.tail);
120 }
121
122 ~rbtree()
123 {
124 dec();
125 }
126
127 void inc() const
128 {
129 root->inc();
130 tail->inc();
131 }
132
133 void dec() const
134 {
135 traverse(dec_visitor());
136 }
137
138 auto tail_size() const
139 {
140 return size ? ((size - 1) & mask<BL>) + 1 : 0;
141 }
142
143 auto tail_offset() const
144 {
145 return size ? (size - 1) & ~mask<BL> : 0;
146 }
147
148 template <typename Visitor, typename... Args>
149 void traverse(Visitor v, Args&&... args) const
150 {
151 auto tail_off = tail_offset();
152 auto tail_size = size - tail_off;
153
154 if (tail_off) make_regular_sub_pos(root, shift, tail_off).visit(v, args...);
155 else make_empty_regular_pos(root).visit(v, args...);
156
157 make_leaf_sub_pos(tail, tail_size).visit(v, args...);
158 }
159
160 template <typename Visitor, typename... Args>
161 void traverse(Visitor v, size_t first, size_t last, Args&&... args) const
162 {
163 auto tail_off = tail_offset();
164 auto tail_size = size - tail_off;
165
166 if (first < tail_off)
167 make_regular_sub_pos(root, shift, tail_off).visit(
168 v,
169 first,
170 last < tail_off ? last : tail_off,
171 args...);
172 if (last > tail_off)
173 make_leaf_sub_pos(tail, tail_size).visit(
174 v,
175 first > tail_off ? first - tail_off : 0,
176 last - tail_off,
177 args...);
178 }
179
180 template <typename Visitor, typename... Args>
181 bool traverse_p(Visitor v, Args&&... args) const
182 {
183 auto tail_off = tail_offset();
184 auto tail_size = size - tail_off;
185 return (tail_off
186 ? make_regular_sub_pos(root, shift, tail_off).visit(v, args...)
187 : make_empty_regular_pos(root).visit(v, args...))
188 && make_leaf_sub_pos(tail, tail_size).visit(v, args...);
189 }
190
191 template <typename Visitor, typename... Args>
192 bool traverse_p(Visitor v, size_t first, size_t last, Args&&... args) const
193 {
194 auto tail_off = tail_offset();
195 auto tail_size = size - tail_off;
196
197 return
198 (first < tail_off
199 ? make_regular_sub_pos(root, shift, tail_off).visit(
200 v,
201 first,
202 last < tail_off ? last : tail_off,
203 args...)
204 : true)
205 && (last > tail_off
206 ? make_leaf_sub_pos(tail, tail_size).visit(
207 v,
208 first > tail_off ? first - tail_off : 0,
209 last - tail_off,
210 args...)
211 : true);
212 }
213
214 template <typename Visitor>
215 decltype(auto) descend(Visitor v, size_t idx) const
216 {
217 auto tail_off = tail_offset();
218 return idx >= tail_off
219 ? make_leaf_descent_pos(tail).visit(v, idx)
220 : visit_regular_descent(root, shift, v, idx);
221 }
222
223 template <typename Fn>
224 void for_each_chunk(Fn&& fn) const
225 {
226 traverse(for_each_chunk_visitor{}, std::forward<Fn>(fn));
227 }
228
229 template <typename Fn>
230 void for_each_chunk(size_t first, size_t last, Fn&& fn) const
231 {
232 traverse(for_each_chunk_i_visitor{}, first, last, std::forward<Fn>(fn));
233 }
234
235 template <typename Fn>
236 bool for_each_chunk_p(Fn&& fn) const
237 {
238 return traverse_p(for_each_chunk_p_visitor{}, std::forward<Fn>(fn));
239 }
240
241 template <typename Fn>
242 bool for_each_chunk_p(size_t first, size_t last, Fn&& fn) const
243 {
244 return traverse_p(for_each_chunk_p_i_visitor{}, first, last, std::forward<Fn>(fn));
245 }
246
247 bool equals(const rbtree& other) const
248 {
249 if (size != other.size) return false;
250 if (size == 0) return true;
251 return (size <= branches<BL>
252 || make_regular_sub_pos(root, shift, tail_offset()).visit(
253 equals_visitor{}, other.root))
254 && make_leaf_sub_pos(tail, tail_size()).visit(
255 equals_visitor{}, other.tail);
256 }
257
258 void ensure_mutable_tail(edit_t e, count_t n)
259 {
260 if (!tail->can_mutate(e)) {
261 auto new_tail = node_t::copy_leaf_e(e, tail, n);
262 dec_leaf(tail, n);
263 tail = new_tail;
264 }
265 }
266
267 void push_back_mut(edit_t e, T value)
268 {
269 auto tail_off = tail_offset();
270 auto ts = size - tail_off;
271 if (ts < branches<BL>) {
272 ensure_mutable_tail(e, ts);
273 new (&tail->leaf()[ts]) T{std::move(value)};
274 } else {
275 auto new_tail = node_t::make_leaf_e(e, std::move(value));
276 try {
277 if (tail_off == size_t{branches<B>} << shift) {
278 auto new_root = node_t::make_inner_e(e);
279 try {
280 auto path = node_t::make_path_e(e, shift, tail);
281 new_root->inner() [0] = root;
282 new_root->inner() [1] = path;
283 root = new_root;
284 tail = new_tail;
285 shift += B;
286 } catch (...) {
287 node_t::delete_inner_e(new_root);
288 throw;
289 }
290 } else if (tail_off) {
291 auto new_root = make_regular_sub_pos(root, shift, tail_off)
292 .visit(push_tail_mut_visitor<node_t>{}, e, tail);
293 root = new_root;
294 tail = new_tail;
295 } else {
296 auto new_root = node_t::make_path_e(e, shift, tail);
297 assert(tail_off == 0);
298 dec_empty_regular(root);
299 root = new_root;
300 tail = new_tail;
301 }
302 } catch (...) {
303 node_t::delete_leaf(new_tail, 1);
304 throw;
305 }
306 }
307 ++size;
308 }
309
310 rbtree push_back(T value) const
311 {
312 auto tail_off = tail_offset();
313 auto ts = size - tail_off;
314 if (ts < branches<BL>) {
315 auto new_tail = node_t::copy_leaf_emplace(tail, ts,
316 std::move(value));
317 return { size + 1, shift, root->inc(), new_tail };
318 } else {
319 auto new_tail = node_t::make_leaf_n(1, std::move(value));
320 try {
321 if (tail_off == size_t{branches<B>} << shift) {
322 auto new_root = node_t::make_inner_n(2);
323 try {
324 auto path = node_t::make_path(shift, tail);
325 new_root->inner() [0] = root;
326 new_root->inner() [1] = path;
327 root->inc();
328 tail->inc();
329 return { size + 1, shift + B, new_root, new_tail };
330 } catch (...) {
331 node_t::delete_inner(new_root, 2);
332 throw;
333 }
334 } else if (tail_off) {
335 auto new_root = make_regular_sub_pos(root, shift, tail_off)
336 .visit(push_tail_visitor<node_t>{}, tail);
337 tail->inc();
338 return { size + 1, shift, new_root, new_tail };
339 } else {
340 auto new_root = node_t::make_path(shift, tail);
341 tail->inc();
342 return { size + 1, shift, new_root, new_tail };
343 }
344 } catch (...) {
345 node_t::delete_leaf(new_tail, 1);
346 throw;
347 }
348 }
349 }
350
351 const T* array_for(size_t index) const
352 {
353 return descend(array_for_visitor<T>(), index);
354 }
355
356 T& get_mut(edit_t e, size_t idx)
357 {
358 auto tail_off = tail_offset();
359 if (idx >= tail_off) {
360 ensure_mutable_tail(e, size - tail_off);
361 return tail->leaf() [idx & mask<BL>];
362 } else {
363 return make_regular_sub_pos(root, shift, tail_off)
364 .visit(get_mut_visitor<node_t>{}, idx, e, &root);
365 }
366 }
367
368 const T& get(size_t index) const
369 {
370 return descend(get_visitor<T>(), index);
371 }
372
373 const T& get_check(size_t index) const
374 {
375 if (index >= size)
376 throw std::out_of_range{"index out of range"};
377 return descend(get_visitor<T>(), index);
378 }
379
380 const T& front() const
381 {
382 return get(0);
383 }
384
385 const T& back() const
386 {
387 return tail->leaf()[(size - 1) & mask<BL>];
388 }
389
390 template <typename FnT>
391 void update_mut(edit_t e, size_t idx, FnT&& fn)
392 {
393 auto& elem = get_mut(e, idx);
394 elem = std::forward<FnT>(fn) (std::move(elem));
395 }
396
397 template <typename FnT>
398 rbtree update(size_t idx, FnT&& fn) const
399 {
400 auto tail_off = tail_offset();
401 if (idx >= tail_off) {
402 auto tail_size = size - tail_off;
403 auto new_tail = make_leaf_sub_pos(tail, tail_size)
404 .visit(update_visitor<node_t>{}, idx - tail_off, fn);
405 return { size, shift, root->inc(), new_tail };
406 } else {
407 auto new_root = make_regular_sub_pos(root, shift, tail_off)
408 .visit(update_visitor<node_t>{}, idx, fn);
409 return { size, shift, new_root, tail->inc() };
410 }
411 }
412
413 void assoc_mut(edit_t e, size_t idx, T value)
414 {
415 update_mut(e, idx, [&] (auto&&) {
416 return std::move(value);
417 });
418 }
419
420 rbtree assoc(size_t idx, T value) const
421 {
422 return update(idx, [&] (auto&&) {
423 return std::move(value);
424 });
425 }
426
427 rbtree take(size_t new_size) const
428 {
429 auto tail_off = tail_offset();
430 if (new_size == 0) {
431 return empty();
432 } else if (new_size >= size) {
433 return *this;
434 } else if (new_size > tail_off) {
435 auto new_tail = node_t::copy_leaf(tail, new_size - tail_off);
436 return { new_size, shift, root->inc(), new_tail };
437 } else {
438 using std::get;
439 auto l = new_size - 1;
440 auto v = slice_right_visitor<node_t>();
441 auto r = make_regular_sub_pos(root, shift, tail_off).visit(v, l);
442 auto new_shift = get<0>(r);
443 auto new_root = get<1>(r);
444 auto new_tail = get<3>(r);
445 if (new_root) {
446 IMMER_ASSERT_TAGGED(new_root->compute_shift() == get<0>(r));
447 assert(new_root->check(new_shift, new_size - get<2>(r)));
448 return { new_size, new_shift, new_root, new_tail };
449 } else {
450 return { new_size, BL, empty().root->inc(), new_tail };
451 }
452 }
453 }
454
455 void take_mut(edit_t e, size_t new_size)
456 {
457 auto tail_off = tail_offset();
458 if (new_size == 0) {
459 // todo: more efficient?
460 *this = empty();
461 } else if (new_size >= size) {
462 return;
463 } else if (new_size > tail_off) {
464 auto ts = size - tail_off;
465 auto newts = new_size - tail_off;
466 if (tail->can_mutate(e)) {
467 destroy_n(tail->leaf() + newts, ts - newts);
468 } else {
469 auto new_tail = node_t::copy_leaf_e(e, tail, newts);
470 dec_leaf(tail, ts);
471 tail = new_tail;
472 }
473 size = new_size;
474 return;
475 } else {
476 using std::get;
477 auto l = new_size - 1;
478 auto v = slice_right_mut_visitor<node_t>();
479 auto r = make_regular_sub_pos(root, shift, tail_off).visit(v, l, e);
480 auto new_shift = get<0>(r);
481 auto new_root = get<1>(r);
482 auto new_tail = get<3>(r);
483 if (new_root) {
484 root = new_root;
485 shift = new_shift;
486 } else {
487 root = empty().root->inc();
488 shift = BL;
489 }
490 dec_leaf(tail, size - tail_off);
491 size = new_size;
492 tail = new_tail;
493 return;
494 }
495 }
496
497 bool check_tree() const
498 {
499#if IMMER_DEBUG_DEEP_CHECK
500 assert(shift >= BL);
501 assert(tail_offset() <= size);
502 assert(check_root());
503 assert(check_tail());
504#endif
505 return true;
506 }
507
508 bool check_tail() const
509 {
510#if IMMER_DEBUG_DEEP_CHECK
511 if (tail_size() > 0)
512 assert(tail->check(0, tail_size()));
513#endif
514 return true;
515 }
516
517 bool check_root() const
518 {
519#if IMMER_DEBUG_DEEP_CHECK
520 if (tail_offset() > 0)
521 assert(root->check(shift, tail_offset()));
522 else {
523 IMMER_ASSERT_TAGGED(root->kind() == node_t::kind_t::inner);
524 assert(shift == BL);
525 }
526#endif
527 return true;
528 }
529};
530
531} // namespace rbts
532} // namespace detail
533} // namespace immer
534