Pixie
Loading...
Searching...
No Matches
btree.h
1#pragma once
2
3#include <pixie/bits.h>
4#include <pixie/memory_usage.h>
5#include <pixie/rank_select/support.h>
6#include <pixie/rmm.h>
7
8#include <algorithm>
9#include <array>
10#include <bit>
11#include <cstddef>
12#include <cstdint>
13#include <limits>
14#include <optional>
15#include <span>
16#include <stdexcept>
17#include <type_traits>
18#include <utility>
19#include <vector>
20
21namespace pixie::experimental {
22
82template <std::size_t HighCacheLines = 4, std::size_t LowFanout = 32>
83class RmMBTree : public RmMBase<RmMBTree<HighCacheLines, LowFanout>> {
84 public:
85 static_assert(HighCacheLines > 0);
86 static_assert(LowFanout > 0);
87
88 static constexpr std::size_t npos =
90 static constexpr std::size_t kBlockBits = 512;
91 static constexpr std::size_t kBlockWords = kBlockBits / 64;
92 static constexpr std::size_t kCacheLineBytes = 64;
93 static constexpr std::size_t kLowFanout = LowFanout;
94 static constexpr std::size_t kHighFanout =
95 std::max<std::size_t>(2, (512 * HighCacheLines) / (4 * 64));
96 static constexpr std::size_t kMaxFanout = std::max(kLowFanout, kHighFanout);
97 static_assert(kMaxFanout <= 64);
98 static_assert(
99 kBlockBits * kLowFanout <=
100 static_cast<std::size_t>(std::numeric_limits<std::int16_t>::max()));
101
102 RmMBTree() = default;
103 RmMBTree(const RmMBTree&) = default;
104 RmMBTree(RmMBTree&&) noexcept = default;
105 RmMBTree& operator=(const RmMBTree&) = default;
106 RmMBTree& operator=(RmMBTree&&) noexcept = default;
107
112 std::size_t position = npos;
113 int value = 0;
114 };
115
126 explicit RmMBTree(std::span<const std::uint64_t> words,
127 std::size_t bit_count,
128 std::size_t = kBlockBits) {
129 build(words, bit_count);
130 }
131
140 RmMBTree(std::span<const std::uint64_t> words,
141 std::size_t bit_count,
143 std::optional<std::size_t> one_count = std::nullopt,
144 std::size_t = kBlockBits) {
145 build(words, bit_count, select_support, one_count);
146 }
147
148 std::size_t size_impl() const { return bit_count_; }
149
150 std::size_t rank1_impl(std::size_t end_position) const {
151 return rank_index_ ? rank_index_->rank(end_position) : 0;
152 }
153
154 std::size_t rank0_impl(std::size_t end_position) const {
155 return rank_index_ ? rank_index_->rank0(end_position) : 0;
156 }
157
158 std::size_t select1_impl(std::size_t rank) const {
159 if (!rank_index_ || rank == 0) {
160 return npos;
161 }
162 const std::size_t position = rank_index_->select(rank);
163 return position < bit_count_ ? position : npos;
164 }
165
166 std::size_t select0_impl(std::size_t rank) const {
167 if (!rank_index_ || rank == 0) {
168 return npos;
169 }
170 const std::size_t position = rank_index_->select0(rank);
171 return position < bit_count_ ? position : npos;
172 }
173
174 std::size_t rank10_impl(std::size_t end_position) const {
175 if (end_position <= 1 || bit_count_ == 0) {
176 return 0;
177 }
178 end_position = std::min(end_position, bit_count_);
179 std::size_t count = 0;
180 for (std::size_t position = 0; position + 1 < end_position; ++position) {
181 count += bit(position) == 1 && bit(position + 1) == 0;
182 }
183 return count;
184 }
185
186 std::size_t select10_impl(std::size_t rank) const {
187 if (rank == 0) {
188 return npos;
189 }
190 for (std::size_t position = 0; position + 1 < bit_count_; ++position) {
191 if (bit(position) == 1 && bit(position + 1) == 0 && --rank == 0) {
192 return position;
193 }
194 }
195 return npos;
196 }
197
198 int excess_impl(std::size_t end_position) const {
199 end_position = std::min(end_position, bit_count_);
200 return static_cast<int>(
201 2 * static_cast<std::int64_t>(rank1_impl(end_position)) -
202 static_cast<std::int64_t>(end_position));
203 }
204
205 std::size_t fwdsearch_impl(std::size_t start_position, int delta) const {
206 if (start_position >= bit_count_) {
207 return npos;
208 }
209
210 const std::size_t block_index = start_position / kBlockBits;
211 const std::size_t block_begin = block_index * kBlockBits;
212 const std::size_t start_offset = start_position - block_begin;
213 const std::size_t block_result =
214 find_fwd_in_block(block_index, start_offset, delta);
215 if (block_result != npos) {
216 return block_result;
217 }
218
219 const std::int64_t relative_target =
220 block_excess_at_local(block_index, start_offset) + delta;
221 const std::int64_t block_start_excess = prefix_excess_impl(block_begin);
222 const std::int64_t target = block_start_excess + relative_target;
223 std::size_t level = 0;
224 std::size_t index = block_index;
225 while (has_parent_level(level)) {
226 const std::size_t fanout = fanout_to_parent(level);
227 const std::size_t parent = index / fanout;
228 const std::size_t sibling_end =
229 std::min(level_count(level), parent * fanout + fanout);
230 const NodeScanResult scan = scan_children_fwd(
231 level, parent, index % fanout + 1, sibling_end - parent * fanout,
232 target, prefix_excess_impl(node_end_bit(level, index)));
233 if (scan.found) {
234 const std::size_t result =
235 descend_fwd(level, scan.index, target, scan.node_start_excess);
236 if (result != npos) {
237 return result;
238 }
239 }
240 level += 1;
241 index = parent;
242 }
243 return npos;
244 }
245
246 std::size_t bwdsearch_impl(std::size_t start_position, int delta) const {
247 if (start_position == 0 || start_position > bit_count_) {
248 return npos;
249 }
250
251 const std::size_t block_index = (start_position - 1) / kBlockBits;
252 const std::size_t block_begin = block_index * kBlockBits;
253 const std::size_t end_offset = start_position - block_begin;
254 const std::size_t block_result =
255 find_bwd_in_block(block_index, end_offset, delta);
256 if (block_result != npos) {
257 return block_result;
258 }
259
260 const std::int64_t relative_target =
261 block_excess_at_local(block_index, end_offset) + delta;
262 const std::int64_t block_start_excess = prefix_excess_impl(block_begin);
263 const std::int64_t target = block_start_excess + relative_target;
264 std::size_t level = 0;
265 std::size_t index = block_index;
266 while (has_parent_level(level)) {
267 const std::size_t fanout = fanout_to_parent(level);
268 const std::size_t parent = index / fanout;
269 const NodeScanResult scan =
270 scan_children_bwd(level, parent, 0, index % fanout, target,
271 prefix_excess_impl(node_start_bit(level, index)));
272 if (scan.found) {
273 if (scan.boundary_only) {
274 return node_start_bit(level, scan.index);
275 }
276 const std::size_t result =
277 descend_bwd(level, scan.index, target, scan.node_start_excess);
278 if (result != npos) {
279 return result;
280 }
281 }
282 level += 1;
283 index = parent;
284 }
285 return npos;
286 }
287
288 std::size_t range_min_query_pos_impl(std::size_t range_begin,
289 std::size_t range_end) const {
290 if (range_begin > range_end || range_end >= bit_count_) {
291 return npos;
292 }
293 return range_extreme_query_pos(range_begin, range_end, true);
294 }
295
296 int range_min_query_val_impl(std::size_t range_begin,
297 std::size_t range_end) const {
298 if (range_begin > range_end || range_end >= bit_count_) {
299 return 0;
300 }
301 return range_extreme_query_val(range_begin, range_end, true);
302 }
303
310 RangeMinQueryResult range_min_query_result(std::size_t range_begin,
311 std::size_t range_end) const {
312 if (range_begin > range_end || range_end >= bit_count_) {
313 return {};
314 }
315 const auto result = range_extreme_query(range_begin, range_end, true);
316 return {result.position, static_cast<int>(result.value)};
317 }
318
319 std::size_t range_max_query_pos_impl(std::size_t range_begin,
320 std::size_t range_end) const {
321 if (range_begin > range_end || range_end >= bit_count_) {
322 return npos;
323 }
324 return range_extreme_query_pos(range_begin, range_end, false);
325 }
326
327 int range_max_query_val_impl(std::size_t range_begin,
328 std::size_t range_end) const {
329 if (range_begin > range_end || range_end >= bit_count_) {
330 return 0;
331 }
332 return range_extreme_query_val(range_begin, range_end, false);
333 }
334
342 std::size_t memory_usage_bytes() const {
343 std::size_t bytes = sizeof(*this);
344 bytes += pixie::optional_nested_owned_memory_bytes(rank_index_);
345 bytes += pixie::vector_capacity_bytes(level_counts_);
346 bytes += pixie::vector_capacity_bytes(low_levels_);
347 bytes += pixie::vector_capacity_bytes(high_levels_);
348 for (const std::vector<LowNode>& level : low_levels_) {
349 bytes += pixie::vector_capacity_bytes(level);
350 }
351 for (const std::vector<HighNode>& level : high_levels_) {
352 bytes += pixie::vector_capacity_bytes(level);
353 }
354 return bytes;
355 }
356
357 std::size_t mincount_impl(std::size_t range_begin,
358 std::size_t range_end) const {
359 if (range_begin > range_end || range_end >= bit_count_) {
360 return 0;
361 }
362 return range_min_stats(range_begin, range_end).count;
363 }
364
365 std::size_t minselect_impl(std::size_t range_begin,
366 std::size_t range_end,
367 std::size_t rank) const {
368 if (range_begin > range_end || range_end >= bit_count_ || rank == 0) {
369 return npos;
370 }
371 const RangeMinStats stats = range_min_stats(range_begin, range_end);
372 if (rank > stats.count) {
373 return npos;
374 }
375 return range_min_select(range_begin, range_end, stats.value, rank);
376 }
377
378 std::size_t close_impl(std::size_t open_position) const {
379 if (open_position >= bit_count_) {
380 return npos;
381 }
382 if (!bit(open_position)) {
383 return open_position;
384 }
385 return fwd_excess_at(open_position, -1);
386 }
387
388 std::size_t open_impl(std::size_t close_position) const {
389 if (close_position >= bit_count_) {
390 return npos;
391 }
392 if (bit(close_position)) {
393 return close_position;
394 }
395 return bwdsearch_impl(close_position + 1, 0);
396 }
397
398 std::size_t enclose_impl(std::size_t position) const {
399 if (position >= bit_count_) {
400 return npos;
401 }
402 if (!bit(position)) {
403 return open_impl(position);
404 }
405 return bwdsearch_impl(position + 1, -2);
406 }
407
408 private:
420 std::size_t fwd_excess_at(std::size_t position, int delta) const {
421 if (position >= bit_count_) {
422 return npos;
423 }
424 if (position + 1 >= bit_count_) {
425 return npos;
426 }
427 return fwdsearch_impl(position + 1, delta);
428 }
429
430 struct Summary {
431 std::uint64_t size_bits = 0;
432 std::uint64_t ones = 0;
433 std::int64_t block_excess = 0;
434 std::int64_t min_excess = 0;
435 std::int64_t max_excess = 0;
436 std::uint64_t min_count = 0;
437 };
438
439 template <class Excess, class Count, std::size_t Fanout>
440 struct alignas(kCacheLineBytes) SummaryNode {
441 using ExcessType = Excess;
442 using CountType = Count;
443 static constexpr std::size_t kFanout = Fanout;
444 std::array<Excess, Fanout> prefix_excess{};
445 std::array<Excess, Fanout> min_excess{};
446 std::array<Excess, Fanout> max_excess{};
447 std::array<Count, Fanout> min_count{};
448 };
449
450 using LowNode = SummaryNode<std::int16_t, std::uint16_t, kLowFanout>;
451 using HighNode = SummaryNode<std::int64_t, std::uint64_t, kHighFanout>;
452 static_assert(alignof(LowNode) == kCacheLineBytes);
453 static_assert(alignof(HighNode) == kCacheLineBytes);
454 static_assert(sizeof(LowNode) % kCacheLineBytes == 0);
455 static_assert(sizeof(HighNode) % kCacheLineBytes == 0);
456
457 struct ByteAgg {
458 std::int8_t block_excess = 0;
459 std::int8_t min_excess = 0;
460 std::int8_t max_excess = 0;
461 std::uint8_t min_count = 0;
462 std::uint8_t pos_first_min = 0;
463 std::uint8_t pos_first_max = 0;
464 };
465
466 static constexpr std::size_t kSearchChunkBits = 128;
467 static constexpr std::size_t kSearchChunkWords = kSearchChunkBits / 64;
468 static constexpr std::size_t kSearchChunkCount =
469 kBlockBits / kSearchChunkBits;
470
478 static const std::array<ByteAgg, 256>& byte_lut() {
479 static const std::array<ByteAgg, 256> table = [] {
480 std::array<ByteAgg, 256> result{};
481 for (int byte_value = 0; byte_value < 256; ++byte_value) {
482 ByteAgg agg;
483 int current = 0;
484 int minimum = std::numeric_limits<int>::max();
485 int maximum = std::numeric_limits<int>::min();
486 const auto bit_at = [&](int bit_index) {
487 return (byte_value >> bit_index) & 1;
488 };
489 for (int bit_index = 0; bit_index < 8; ++bit_index) {
490 const int value = bit_at(bit_index);
491 current += value ? 1 : -1;
492 if (current < minimum) {
493 minimum = current;
494 agg.min_count = 1;
495 agg.pos_first_min = static_cast<std::uint8_t>(bit_index);
496 } else if (current == minimum) {
497 ++agg.min_count;
498 }
499 if (current > maximum) {
500 maximum = current;
501 agg.pos_first_max = static_cast<std::uint8_t>(bit_index);
502 }
503 }
504 agg.block_excess = static_cast<std::int8_t>(current);
505 agg.min_excess = static_cast<std::int8_t>(minimum);
506 agg.max_excess = static_cast<std::int8_t>(maximum);
507 result[static_cast<std::size_t>(byte_value)] = agg;
508 }
509 return result;
510 }();
511 return table;
512 }
513
524 void build(std::span<const std::uint64_t> words, std::size_t bit_count) {
525 build(words, bit_count, RankSelectSupport<>::SelectSupport::kBoth,
526 std::nullopt);
527 }
528
532 void build(std::span<const std::uint64_t> words,
533 std::size_t bit_count,
534 RankSelectSupport<>::SelectSupport select_support,
535 std::optional<std::size_t> one_count = std::nullopt) {
536 const std::size_t required_words = (bit_count + 63) / 64;
537 if (words.size() < required_words) {
538 throw std::invalid_argument(
539 "RmMBTree input span is shorter than bit_count");
540 }
541
542 bits_ = words;
543 bit_count_ = bit_count;
544 rank_index_.emplace(words, bit_count, select_support, one_count);
545 block_count_ = (bit_count_ + kBlockBits - 1) / kBlockBits;
546 std::vector<Summary> block_summaries(block_count_);
547
548 for (std::size_t block_index = 0; block_index < block_count_;
549 ++block_index) {
550 const std::size_t block_begin = block_index * kBlockBits;
551 const std::size_t block_end =
552 std::min(bit_count_, block_begin + kBlockBits);
553 block_summaries[block_index] =
554 summarize_bits(block_begin, block_end - block_begin);
555 }
556
557 build_levels(block_summaries);
558 }
559
567 void build_levels(const std::vector<Summary>& block_summaries) {
568 level_counts_.clear();
569 low_levels_.clear();
570 high_levels_.clear();
571 top_summary_ = Summary{};
572 level_counts_.push_back(block_summaries.size());
573 if (block_summaries.empty()) {
574 return;
575 }
576
577 std::vector<Summary> current = block_summaries;
578 current = build_parent_level<LowNode>(current, low_levels_.emplace_back());
579 level_counts_.push_back(current.size());
580 current =
581 build_parent_level<HighNode>(current, high_levels_.emplace_back());
582 level_counts_.push_back(current.size());
583 while (current.size() > 1) {
584 current =
585 build_parent_level<HighNode>(current, high_levels_.emplace_back());
586 level_counts_.push_back(current.size());
587 }
588 top_summary_ = current.front();
589 }
590
601 template <class Node>
602 static std::vector<Summary> build_parent_level(const std::vector<Summary>& in,
603 std::vector<Node>& nodes) {
604 constexpr std::size_t fanout = Node::kFanout;
605 std::vector<Summary> out((in.size() + fanout - 1) / fanout);
606 nodes.resize(out.size());
607 for (std::size_t parent = 0; parent < out.size(); ++parent) {
608 const std::size_t begin = parent * fanout;
609 const std::size_t end = std::min(in.size(), begin + fanout);
610 Summary combined;
611 for (std::size_t i = begin; i < end; ++i) {
612 store_child_summary(nodes[parent], i - begin, combined.block_excess,
613 in[i]);
614 combined = append(combined, in[i]);
615 }
616 out[parent] = combined;
617 }
618 return out;
619 }
620
633 template <class Node>
634 static void store_child_summary(Node& node,
635 std::size_t slot,
636 std::int64_t prefix_excess,
637 const Summary& summary) {
638 node.prefix_excess[slot] =
639 static_cast<typename decltype(node.prefix_excess)::value_type>(
640 prefix_excess + summary.block_excess);
641 node.min_excess[slot] =
642 static_cast<typename decltype(node.min_excess)::value_type>(
643 summary.min_excess);
644 node.max_excess[slot] =
645 static_cast<typename decltype(node.max_excess)::value_type>(
646 summary.max_excess);
647 node.min_count[slot] =
648 static_cast<typename decltype(node.min_count)::value_type>(
649 summary.min_count);
650 }
651
661 Summary summarize_bits(std::size_t begin, std::size_t length) const {
662 if (length == kBlockBits && (begin % kBlockBits) == 0) {
663 const std::size_t block_index = begin / kBlockBits;
664 if (full_block_has_words(block_index)) {
665 return summarize_full_block(block_index);
666 }
667 }
668
669 Summary summary;
670 summary.size_bits = length;
671 if (length == 0) {
672 return summary;
673 }
674 int current = 0;
675 int minimum = std::numeric_limits<int>::max();
676 int maximum = std::numeric_limits<int>::min();
677 for (std::size_t offset = 0; offset < length; ++offset) {
678 const std::uint8_t value = bit(begin + offset);
679 summary.ones += value;
680 current += value ? 1 : -1;
681 if (current < minimum) {
682 minimum = current;
683 summary.min_count = 1;
684 } else if (current == minimum) {
685 ++summary.min_count;
686 }
687 if (current > maximum) {
688 maximum = current;
689 }
690 }
691 summary.block_excess = current;
692 summary.min_excess = minimum;
693 summary.max_excess = maximum;
694 return summary;
695 }
696
703 Summary summarize_full_block(std::size_t block_index) const {
704 const std::uint64_t* block = bits_.data() + block_index * kBlockWords;
705 Summary summary;
706 for (std::size_t chunk = 0; chunk < kSearchChunkCount; ++chunk) {
707 summary = append(summary,
708 summarize_128_chunk(block + chunk * kSearchChunkWords));
709 }
710 return summary;
711 }
712
719 static Summary summarize_128_chunk(const std::uint64_t* chunk) {
720 Summary summary;
721 summary.size_bits = kSearchChunkBits;
722 summary.ones = std::popcount(chunk[0]) + std::popcount(chunk[1]);
723 summary.block_excess = 2 * static_cast<std::int64_t>(summary.ones) -
724 static_cast<std::int64_t>(kSearchChunkBits);
725
726 const ExcessResult minimum = excess_min_128(chunk, 1, kSearchChunkBits);
727 summary.min_excess = minimum.min_excess;
728
729 const std::array<std::uint64_t, kSearchChunkWords> inverted = {~chunk[0],
730 ~chunk[1]};
731 const ExcessResult inverted_min =
732 excess_min_128(inverted.data(), 1, kSearchChunkBits);
733 summary.max_excess = -static_cast<std::int64_t>(inverted_min.min_excess);
734
735 std::uint64_t minimum_positions[kSearchChunkWords];
736 excess_positions_128(chunk, static_cast<int>(summary.min_excess),
737 minimum_positions);
738 summary.min_count = std::popcount(minimum_positions[0]) +
739 std::popcount(minimum_positions[1]);
740 return summary;
741 }
742
752 static Summary append(Summary left, const Summary& right) {
753 if (left.size_bits == 0) {
754 return right;
755 }
756 if (right.size_bits == 0) {
757 return left;
758 }
759 Summary result;
760 result.size_bits = left.size_bits + right.size_bits;
761 result.ones = left.ones + right.ones;
762 result.block_excess = left.block_excess + right.block_excess;
763 result.min_excess =
764 std::min(left.min_excess, left.block_excess + right.min_excess);
765 result.max_excess =
766 std::max(left.max_excess, left.block_excess + right.max_excess);
767 result.min_count = 0;
768 if (left.min_excess == result.min_excess) {
769 result.min_count += left.min_count;
770 }
771 if (left.block_excess + right.min_excess == result.min_excess) {
772 result.min_count += right.min_count;
773 }
774 return result;
775 }
776
784 std::size_t block_size(std::size_t block_index) const {
785 const std::size_t begin = block_index * kBlockBits;
786 return std::min(bit_count_ - begin, kBlockBits);
787 }
788
796 bool full_block_has_words(std::size_t block_index) const {
797 return block_size(block_index) == kBlockBits &&
798 (block_index + 1) * kBlockWords <= bits_.size();
799 }
800
810 std::int64_t block_excess_at_local(std::size_t block_index,
811 std::size_t offset) const {
812 const std::size_t length = block_size(block_index);
813 offset = std::min(offset, length);
814 if (offset == 0) {
815 return 0;
816 }
817
818 const std::size_t first_word = block_index * kBlockWords;
819 std::size_t remaining = offset;
820 std::int64_t ones = 0;
821 std::size_t word_offset = 0;
822 while (remaining >= 64) {
823 ones += std::popcount(bits_[first_word + word_offset]);
824 remaining -= 64;
825 ++word_offset;
826 }
827 if (remaining != 0) {
828 ones += std::popcount(bits_[first_word + word_offset] &
829 first_bits_mask(remaining));
830 }
831 return 2 * ones - static_cast<std::int64_t>(offset);
832 }
833
840 static int chunk_excess_128(const std::uint64_t* chunk) {
841 return 2 * static_cast<int>(std::popcount(chunk[0]) +
842 std::popcount(chunk[1])) -
843 static_cast<int>(kSearchChunkBits);
844 }
845
857 std::size_t find_fwd_in_block(std::size_t block_index,
858 std::size_t start_offset,
859 std::int64_t delta) const {
860 const std::size_t length = block_size(block_index);
861 if (start_offset >= length) {
862 return npos;
863 }
864
865 if (full_block_has_words(block_index)) {
866 const std::uint64_t* block = bits_.data() + block_index * kBlockWords;
867 const std::size_t first_chunk = start_offset / kSearchChunkBits;
868 std::int64_t target = delta;
869 for (std::size_t chunk = first_chunk; chunk < kSearchChunkCount;
870 ++chunk) {
871 const std::uint64_t* chunk_words = block + chunk * kSearchChunkWords;
872 const std::size_t local_start =
873 chunk == first_chunk ? start_offset - chunk * kSearchChunkBits : 0;
874 if (chunk == first_chunk) {
875 target += prefix_excess_128(chunk_words, local_start);
876 }
877 int block_excess = 0;
878 const std::size_t offset = forward_search_128(
879 chunk_words, static_cast<int>(target), local_start, &block_excess);
880 if (offset != kSearchChunkBits) {
881 return block_index * kBlockBits + chunk * kSearchChunkBits + offset;
882 }
883 target -= block_excess;
884 }
885 return npos;
886 }
887
888 std::int64_t current = block_excess_at_local(block_index, start_offset);
889 const std::int64_t relative_target = current + delta;
890 const std::size_t block_begin = block_index * kBlockBits;
891 for (std::size_t offset = start_offset; offset < length; ++offset) {
892 current += bit(block_begin + offset) ? 1 : -1;
893 if (current == relative_target) {
894 return block_index * kBlockBits + offset;
895 }
896 }
897 return npos;
898 }
899
911 std::size_t find_bwd_in_block(std::size_t block_index,
912 std::size_t end_offset,
913 std::int64_t delta) const {
914 if (end_offset == 0) {
915 return npos;
916 }
917 const std::size_t max_prefix_length = end_offset - 1;
918
919 if (full_block_has_words(block_index)) {
920 const std::uint64_t* block = bits_.data() + block_index * kBlockWords;
921 const std::size_t last_chunk = max_prefix_length / kSearchChunkBits;
922 std::int64_t target = delta;
923 for (std::size_t chunk = last_chunk + 1; chunk > 0;) {
924 --chunk;
925 const std::uint64_t* chunk_words = block + chunk * kSearchChunkWords;
926 const std::size_t local_end =
927 chunk == last_chunk ? end_offset - chunk * kSearchChunkBits
928 : kSearchChunkBits;
929 if (chunk == last_chunk) {
930 target += prefix_excess_128(chunk_words, local_end);
931 }
932 int block_excess = 0;
933 const std::size_t offset = backward_search_128(
934 chunk_words, static_cast<int>(target), local_end, &block_excess);
935 if (offset != kSearchChunkBits) {
936 return block_index * kBlockBits + chunk * kSearchChunkBits + offset;
937 }
938 if (chunk > 0) {
939 target += chunk_excess_128(block + (chunk - 1) * kSearchChunkWords);
940 }
941 }
942 return npos;
943 }
944
945 const std::int64_t relative_target =
946 block_excess_at_local(block_index, end_offset) + delta;
947 std::int64_t current =
948 block_excess_at_local(block_index, max_prefix_length);
949 const std::size_t block_begin = block_index * kBlockBits;
950 for (std::size_t prefix_length = max_prefix_length; prefix_length > 0;
951 --prefix_length) {
952 if (current == relative_target) {
953 return block_index * kBlockBits + prefix_length;
954 }
955 current -= bit(block_begin + prefix_length - 1) ? 1 : -1;
956 }
957 return relative_target == 0 ? block_index * kBlockBits : npos;
958 }
959
971 std::size_t descend_fwd(std::size_t level,
972 std::size_t index,
973 std::int64_t target,
974 std::int64_t node_start_excess) const {
975 while (level > 0) {
976 const std::size_t child_level = level - 1;
977 const std::size_t fanout = fanout_to_parent(child_level);
978 const std::size_t child_begin = index * fanout;
979 const std::size_t child_end =
980 std::min(level_count(child_level), child_begin + fanout);
981 const NodeScanResult scan =
982 scan_children_fwd(child_level, index, 0, child_end - child_begin,
983 target, node_start_excess);
984 if (!scan.found) {
985 return npos;
986 }
987 index = scan.index;
988 level = child_level;
989 node_start_excess = scan.node_start_excess;
990 }
991 return find_fwd_in_block(index, 0, target - node_start_excess);
992 }
993
1005 std::size_t descend_bwd(std::size_t level,
1006 std::size_t index,
1007 std::int64_t target,
1008 std::int64_t node_start_excess) const {
1009 while (level > 0) {
1010 const std::size_t child_level = level - 1;
1011 const std::size_t fanout = fanout_to_parent(child_level);
1012 const std::size_t child_begin = index * fanout;
1013 const std::size_t child_end =
1014 std::min(level_count(child_level), child_begin + fanout);
1015 std::int64_t child_start_excess =
1016 node_start_excess + summary_at(level, index).block_excess;
1017 const NodeScanResult scan =
1018 scan_children_bwd(child_level, index, 0, child_end - child_begin,
1019 target, child_start_excess);
1020 if (!scan.found) {
1021 return npos;
1022 }
1023 if (scan.boundary_only) {
1024 return node_start_bit(child_level, scan.index);
1025 }
1026 index = scan.index;
1027 level = child_level;
1028 node_start_excess = scan.node_start_excess;
1029 }
1030 const std::int64_t block_excess = summary_at(0, index).block_excess;
1031 return find_bwd_in_block(index, block_size(index),
1032 target - node_start_excess - block_excess);
1033 }
1034
1035 struct NodeScanResult {
1036 bool found = false;
1037 bool boundary_only = false;
1038 std::size_t index = 0;
1039 std::int64_t node_start_excess = 0;
1040 };
1041
1054 NodeScanResult scan_children_fwd(std::size_t child_level,
1055 std::size_t parent,
1056 std::size_t begin_slot,
1057 std::size_t end_slot,
1058 std::int64_t target,
1059 std::int64_t begin_excess) const {
1060 if (begin_slot >= end_slot) {
1061 return {};
1062 }
1063 if (child_level == 0) {
1064 return scan_node_fwd(low_levels_[0][parent], child_level, parent,
1065 begin_slot, end_slot, target, begin_excess);
1066 }
1067 return scan_node_fwd(high_levels_[child_level - 1][parent], child_level,
1068 parent, begin_slot, end_slot, target, begin_excess);
1069 }
1070
1084 NodeScanResult scan_children_bwd(std::size_t child_level,
1085 std::size_t parent,
1086 std::size_t begin_slot,
1087 std::size_t end_slot,
1088 std::int64_t target,
1089 std::int64_t end_excess) const {
1090 if (begin_slot >= end_slot) {
1091 return {};
1092 }
1093 if (child_level == 0) {
1094 return scan_node_bwd(low_levels_[0][parent], child_level, parent,
1095 begin_slot, end_slot, target, end_excess);
1096 }
1097 return scan_node_bwd(high_levels_[child_level - 1][parent], child_level,
1098 parent, begin_slot, end_slot, target, end_excess);
1099 }
1100
1115 template <class Node>
1116 NodeScanResult scan_node_fwd(const Node& node,
1117 std::size_t child_level,
1118 std::size_t parent,
1119 std::size_t begin_slot,
1120 std::size_t end_slot,
1121 std::int64_t target,
1122 std::int64_t begin_excess) const {
1123 const std::int64_t node_base_excess =
1124 begin_excess - prefix_excess_at(node, begin_slot);
1125 for (std::size_t slot = begin_slot; slot < end_slot;) {
1126 const std::size_t lane_count =
1127 std::min(vector_lane_count<Node>(), end_slot - slot);
1128 const std::uint32_t mask = matching_chunk_mask(
1129 node, slot, lane_count, target - node_base_excess, false);
1130 if (mask != 0) {
1131 const std::size_t lane = std::countr_zero(mask);
1132 const std::size_t matched_slot = slot + lane;
1133 return {true, false,
1134 parent * fanout_to_parent(child_level) + matched_slot,
1135 child_start_excess(node, node_base_excess, matched_slot)};
1136 }
1137 slot += lane_count;
1138 }
1139 return {};
1140 }
1141
1156 template <class Node>
1157 NodeScanResult scan_node_bwd(const Node& node,
1158 std::size_t child_level,
1159 std::size_t parent,
1160 std::size_t begin_slot,
1161 std::size_t end_slot,
1162 std::int64_t target,
1163 std::int64_t end_excess) const {
1164 const std::int64_t node_base_excess =
1165 end_excess - prefix_excess_at(node, end_slot);
1166 for (std::size_t slot_end = end_slot; slot_end > begin_slot;) {
1167 const std::size_t lane_count =
1168 std::min(vector_lane_count<Node>(), slot_end - begin_slot);
1169 const std::size_t slot = slot_end - lane_count;
1170 const std::uint32_t mask = matching_chunk_mask(
1171 node, slot, lane_count, target - node_base_excess, true);
1172 if (mask != 0) {
1173 const std::size_t lane =
1174 static_cast<std::size_t>(std::bit_width(mask) - 1);
1175 const std::size_t matched_slot = slot + lane;
1176 const std::int64_t relative_target =
1177 target - child_start_excess(node, node_base_excess, matched_slot);
1178 const bool interior_match =
1179 node.min_excess[matched_slot] <= relative_target &&
1180 relative_target <= node.max_excess[matched_slot];
1181 return {true, !interior_match,
1182 parent * fanout_to_parent(child_level) + matched_slot,
1183 child_start_excess(node, node_base_excess, matched_slot)};
1184 }
1185 slot_end = slot;
1186 }
1187 return {};
1188 }
1189
1200 template <class Node>
1201 static std::int64_t prefix_excess_at(const Node& node, std::size_t slot) {
1202 if (slot == 0) {
1203 return 0;
1204 }
1205 return prefix_through(node, slot - 1);
1206 }
1207
1218 template <class Node>
1219 static std::int64_t child_start_excess(const Node& node,
1220 std::int64_t node_base_excess,
1221 std::size_t slot) {
1222 return node_base_excess + prefix_excess_at(node, slot);
1223 }
1224
1233 template <class Node>
1234 static std::int64_t prefix_through(const Node& node, std::size_t slot) {
1235 return static_cast<std::int64_t>(node.prefix_excess[slot]);
1236 }
1237
1247 template <class Node>
1248 static std::int64_t child_excess(const Node& node, std::size_t slot) {
1249 return prefix_through(node, slot) - prefix_excess_at(node, slot);
1250 }
1251
1258 template <class Node>
1259 static constexpr std::size_t vector_lane_count() {
1260 if constexpr (std::is_same_v<typename Node::ExcessType, std::int16_t>) {
1261 return 16;
1262 } else {
1263 return 4;
1264 }
1265 }
1266
1280 template <class Node>
1281 static std::uint32_t matching_chunk_mask(const Node& node,
1282 std::size_t slot,
1283 std::size_t lane_count,
1284 std::int64_t target_in_node,
1285 bool include_zero_boundary) {
1286#ifdef PIXIE_AVX2_SUPPORT
1287 if constexpr (std::is_same_v<typename Node::ExcessType, std::int16_t>) {
1288 if (lane_count == 16 &&
1289 target_in_node >= std::numeric_limits<std::int16_t>::min() &&
1290 target_in_node <= std::numeric_limits<std::int16_t>::max()) {
1291 alignas(32) std::int16_t prefix_before[16]{};
1292 fill_prefix_before(node, slot, prefix_before);
1293 return rmm_btree_match_mask_i16x16(
1294 prefix_before, node.min_excess.data() + slot,
1295 node.max_excess.data() + slot,
1296 static_cast<std::int16_t>(target_in_node), include_zero_boundary);
1297 }
1298 } else if constexpr (std::is_same_v<typename Node::ExcessType,
1299 std::int64_t>) {
1300 if (lane_count == 4) {
1301 alignas(32) std::int64_t prefix_before[4]{};
1302 fill_prefix_before(node, slot, prefix_before);
1303 return rmm_btree_match_mask_i64x4(
1304 prefix_before, node.min_excess.data() + slot,
1305 node.max_excess.data() + slot, target_in_node,
1306 include_zero_boundary);
1307 }
1308 }
1309#endif
1310 std::uint32_t result = 0;
1311 for (std::size_t lane = 0; lane < lane_count; ++lane) {
1312 const std::int64_t rel =
1313 target_in_node - prefix_excess_at(node, slot + lane);
1314 const bool found = (include_zero_boundary && rel == 0) ||
1315 (node.min_excess[slot + lane] <= rel &&
1316 rel <= node.max_excess[slot + lane]);
1317 if (found) {
1318 result |= std::uint32_t{1} << lane;
1319 }
1320 }
1321 return result;
1322 }
1323
1336 template <class Node>
1337 static void fill_prefix_before(const Node& node,
1338 std::size_t slot,
1339 typename Node::ExcessType* out) {
1340 if (slot == 0) {
1341 out[0] = 0;
1342 for (std::size_t lane = 1; lane < vector_lane_count<Node>(); ++lane) {
1343 out[lane] = node.prefix_excess[lane - 1];
1344 }
1345 return;
1346 }
1347 for (std::size_t lane = 0; lane < vector_lane_count<Node>(); ++lane) {
1348 out[lane] = node.prefix_excess[slot + lane - 1];
1349 }
1350 }
1351
1360 static bool contains_fwd(const Summary& summary,
1361 std::int64_t relative_target) {
1362 return summary.min_excess <= relative_target &&
1363 relative_target <= summary.max_excess;
1364 }
1365
1375 static bool contains_bwd(const Summary& summary,
1376 std::int64_t relative_target) {
1377 return relative_target == 0 || contains_fwd(summary, relative_target);
1378 }
1379
1380 std::int64_t prefix_excess_impl(std::size_t end_position) const {
1381 return 2 * static_cast<std::int64_t>(rank1_impl(end_position)) -
1382 static_cast<std::int64_t>(end_position);
1383 }
1384
1391 bool has_parent_level(std::size_t level) const {
1392 return level + 1 < total_levels() && level_count(level + 1) != 0;
1393 }
1394
1401 std::size_t total_levels() const { return level_counts_.size(); }
1402
1409 std::size_t level_count(std::size_t level) const {
1410 return level < level_counts_.size() ? level_counts_[level] : 0;
1411 }
1412
1421 Summary summary_at(std::size_t level, std::size_t index) const {
1422 if (level + 1 >= total_levels()) {
1423 return top_summary_;
1424 }
1425 const std::size_t parent_level = level + 1;
1426 const std::size_t fanout = fanout_to_parent(level);
1427 const std::size_t parent = index / fanout;
1428 const std::size_t slot = index % fanout;
1429 Summary summary;
1430 if (parent_level == 1) {
1431 const LowNode& node = low_levels_[0][parent];
1432 summary.block_excess = child_excess(node, slot);
1433 summary.min_excess = node.min_excess[slot];
1434 summary.max_excess = node.max_excess[slot];
1435 summary.min_count = node.min_count[slot];
1436 } else {
1437 const HighNode& node = high_levels_[parent_level - 2][parent];
1438 summary.block_excess = child_excess(node, slot);
1439 summary.min_excess = node.min_excess[slot];
1440 summary.max_excess = node.max_excess[slot];
1441 summary.min_count = node.min_count[slot];
1442 }
1443 return summary;
1444 }
1445
1453 static std::size_t fanout_to_parent(std::size_t level) {
1454 return level == 0 ? kLowFanout : kHighFanout;
1455 }
1456
1465 static std::size_t mul_clamped(std::size_t lhs, std::size_t rhs) {
1466 if (lhs != 0 && rhs > std::numeric_limits<std::size_t>::max() / lhs) {
1467 return std::numeric_limits<std::size_t>::max();
1468 }
1469 return lhs * rhs;
1470 }
1471
1479 static std::size_t level_span_bits(std::size_t level) {
1480 std::size_t span = kBlockBits;
1481 if (level >= 1) {
1482 span = mul_clamped(span, kLowFanout);
1483 }
1484 if (level >= 2) {
1485 for (std::size_t i = 2; i <= level; ++i) {
1486 span = mul_clamped(span, kHighFanout);
1487 }
1488 }
1489 return span;
1490 }
1491
1500 std::size_t node_start_bit(std::size_t level, std::size_t index) const {
1501 const std::size_t span = level_span_bits(level);
1502 if (span != 0 && index > std::numeric_limits<std::size_t>::max() / span) {
1503 return bit_count_;
1504 }
1505 return std::min(bit_count_, index * span);
1506 }
1507
1516 std::size_t node_size_bits(std::size_t level, std::size_t index) const {
1517 const std::size_t start = node_start_bit(level, index);
1518 if (start >= bit_count_) {
1519 return 0;
1520 }
1521 return std::min(level_span_bits(level), bit_count_ - start);
1522 }
1523
1532 std::size_t node_end_bit(std::size_t level, std::size_t index) const {
1533 const std::size_t start = node_start_bit(level, index);
1534 const std::uint64_t size = node_size_bits(level, index);
1535 if (size > std::numeric_limits<std::size_t>::max() - start) {
1536 return bit_count_;
1537 }
1538 return std::min(bit_count_, start + static_cast<std::size_t>(size));
1539 }
1540
1541 struct NodeRef {
1542 std::size_t level;
1543 std::size_t index;
1544 };
1545
1546 static constexpr std::size_t kMaxCoverNodes = 512;
1547
1548 struct ScanResult {
1549 std::int64_t block_excess = 0;
1550 std::int64_t min_value = std::numeric_limits<std::int64_t>::max();
1551 std::int64_t max_value = std::numeric_limits<std::int64_t>::min();
1552 std::uint64_t min_count = 0;
1553 std::size_t min_position = npos;
1554 std::size_t max_position = npos;
1555 };
1556
1557 struct MinScanResult {
1558 std::int64_t block_excess = 0;
1559 std::int64_t min_value = std::numeric_limits<std::int64_t>::max();
1560 std::size_t min_position = npos;
1561 };
1562
1563 struct RangeExtremeResult {
1564 std::size_t position = npos;
1565 std::int64_t value = 0;
1566 };
1567
1568 struct RangeMinStats {
1569 std::int64_t value = 0;
1570 std::uint64_t count = 0;
1571 };
1572
1582 std::size_t range_extreme_query_pos(std::size_t range_begin,
1583 std::size_t range_end,
1584 bool find_min) const {
1585 return range_extreme_query(range_begin, range_end, find_min).position;
1586 }
1587
1597 int range_extreme_query_val(std::size_t range_begin,
1598 std::size_t range_end,
1599 bool find_min) const {
1600 return static_cast<int>(
1601 range_extreme_query_value(range_begin, range_end, find_min));
1602 }
1603
1613 std::int64_t range_extreme_query_value(std::size_t range_begin,
1614 std::size_t range_end,
1615 bool find_min) const {
1616 std::int64_t value = 0;
1617 std::int64_t best = find_min ? std::numeric_limits<std::int64_t>::max()
1618 : std::numeric_limits<std::int64_t>::min();
1619
1620 auto consider_value = [&](std::int64_t candidate) {
1621 if ((find_min && candidate < best) || (!find_min && candidate > best)) {
1622 best = candidate;
1623 }
1624 };
1625
1626 const std::size_t range_end_exclusive = range_end + 1;
1627 const std::size_t first_full_block =
1628 (range_begin + kBlockBits - 1) / kBlockBits;
1629 const std::size_t full_begin =
1630 std::min(range_end_exclusive, first_full_block * kBlockBits);
1631 if (range_begin < full_begin) {
1632 if (find_min) {
1633 const MinScanResult scan = scan_min_range(range_begin, full_begin);
1634 consider_value(scan.min_value);
1635 value += scan.block_excess;
1636 } else {
1637 const ScanResult scan = scan_range(range_begin, full_begin);
1638 consider_value(scan.max_value);
1639 value += scan.block_excess;
1640 }
1641 }
1642
1643 const std::size_t last_full_block_exclusive =
1644 range_end_exclusive / kBlockBits;
1645 const std::size_t middle_begin = full_begin;
1646 const std::size_t middle_end =
1647 std::max(middle_begin, last_full_block_exclusive * kBlockBits);
1648 if (middle_begin < middle_end) {
1649 for_each_cover_node(middle_begin, middle_end, [&](NodeRef node) {
1650 const Summary summary = summary_at(node.level, node.index);
1651 consider_value(value +
1652 (find_min ? summary.min_excess : summary.max_excess));
1653 value += summary.block_excess;
1654 return true;
1655 });
1656 }
1657
1658 if (middle_end < range_end_exclusive) {
1659 if (find_min) {
1660 const MinScanResult scan =
1661 scan_min_range(middle_end, range_end_exclusive);
1662 consider_value(value + scan.min_value);
1663 } else {
1664 const ScanResult scan = scan_range(middle_end, range_end_exclusive);
1665 consider_value(value + scan.max_value);
1666 }
1667 }
1668 if (best == std::numeric_limits<std::int64_t>::max() ||
1669 best == std::numeric_limits<std::int64_t>::min()) {
1670 return 0;
1671 }
1672 return best;
1673 }
1674
1685 RangeExtremeResult range_extreme_query(std::size_t range_begin,
1686 std::size_t range_end,
1687 bool find_min) const {
1688 std::int64_t value = 0;
1689 std::int64_t best = find_min ? std::numeric_limits<std::int64_t>::max()
1690 : std::numeric_limits<std::int64_t>::min();
1691 std::size_t best_position = npos;
1692 NodeRef best_node;
1693 std::int64_t prefix_at_best_node = 0;
1694 bool best_is_node = false;
1695
1696 auto consider_point = [&](std::int64_t candidate, std::size_t position) {
1697 if ((find_min && candidate < best) || (!find_min && candidate > best)) {
1698 best = candidate;
1699 best_position = position;
1700 best_is_node = false;
1701 }
1702 };
1703
1704 const std::size_t range_end_exclusive = range_end + 1;
1705 const std::size_t first_full_block =
1706 (range_begin + kBlockBits - 1) / kBlockBits;
1707 const std::size_t full_begin =
1708 std::min(range_end_exclusive, first_full_block * kBlockBits);
1709 if (range_begin < full_begin) {
1710 if (find_min) {
1711 const MinScanResult scan = scan_min_range(range_begin, full_begin);
1712 consider_point(scan.min_value, scan.min_position);
1713 value += scan.block_excess;
1714 } else {
1715 const ScanResult scan = scan_range(range_begin, full_begin);
1716 consider_point(scan.max_value, scan.max_position);
1717 value += scan.block_excess;
1718 }
1719 }
1720
1721 const std::size_t last_full_block_exclusive =
1722 range_end_exclusive / kBlockBits;
1723 const std::size_t middle_begin = full_begin;
1724 const std::size_t middle_end =
1725 std::max(middle_begin, last_full_block_exclusive * kBlockBits);
1726 if (middle_begin < middle_end) {
1727 for_each_cover_node(middle_begin, middle_end, [&](NodeRef node) {
1728 const Summary summary = summary_at(node.level, node.index);
1729 const std::int64_t candidate =
1730 value + (find_min ? summary.min_excess : summary.max_excess);
1731 if ((find_min && candidate < best) || (!find_min && candidate > best)) {
1732 best = candidate;
1733 best_node = node;
1734 prefix_at_best_node = value;
1735 best_is_node = true;
1736 }
1737 value += summary.block_excess;
1738 return true;
1739 });
1740 }
1741
1742 if (middle_end < range_end_exclusive) {
1743 if (find_min) {
1744 const MinScanResult scan =
1745 scan_min_range(middle_end, range_end_exclusive);
1746 consider_point(value + scan.min_value, scan.min_position);
1747 } else {
1748 const ScanResult scan = scan_range(middle_end, range_end_exclusive);
1749 consider_point(value + scan.max_value, scan.max_position);
1750 }
1751 }
1752
1753 if (best_is_node) {
1754 best_position =
1755 descend_first_extreme(best_node.level, best_node.index,
1756 best - prefix_at_best_node, find_min);
1757 }
1758 return {best_position,
1759 best == std::numeric_limits<std::int64_t>::max() ||
1760 best == std::numeric_limits<std::int64_t>::min()
1761 ? 0
1762 : best};
1763 }
1764
1773 RangeMinStats range_min_stats(std::size_t range_begin,
1774 std::size_t range_end) const {
1775 std::int64_t value = 0;
1776 std::int64_t best = std::numeric_limits<std::int64_t>::max();
1777 std::uint64_t count = 0;
1778
1779 auto consider = [&](std::int64_t candidate, std::uint64_t candidate_count) {
1780 if (candidate < best) {
1781 best = candidate;
1782 count = candidate_count;
1783 } else if (candidate == best) {
1784 count += candidate_count;
1785 }
1786 };
1787
1788 const std::size_t range_end_exclusive = range_end + 1;
1789 const std::size_t first_full_block =
1790 (range_begin + kBlockBits - 1) / kBlockBits;
1791 const std::size_t full_begin =
1792 std::min(range_end_exclusive, first_full_block * kBlockBits);
1793 if (range_begin < full_begin) {
1794 const ScanResult scan = scan_range(range_begin, full_begin);
1795 consider(scan.min_value, scan.min_count);
1796 value += scan.block_excess;
1797 }
1798
1799 const std::size_t last_full_block_exclusive =
1800 range_end_exclusive / kBlockBits;
1801 const std::size_t middle_begin = full_begin;
1802 const std::size_t middle_end =
1803 std::max(middle_begin, last_full_block_exclusive * kBlockBits);
1804 if (middle_begin < middle_end) {
1805 for_each_cover_node(middle_begin, middle_end, [&](NodeRef node) {
1806 const Summary summary = summary_at(node.level, node.index);
1807 consider(value + summary.min_excess, summary.min_count);
1808 value += summary.block_excess;
1809 return true;
1810 });
1811 }
1812
1813 if (middle_end < range_end_exclusive) {
1814 const ScanResult scan = scan_range(middle_end, range_end_exclusive);
1815 consider(value + scan.min_value, scan.min_count);
1816 }
1817 return {best == std::numeric_limits<std::int64_t>::max() ? 0 : best, count};
1818 }
1819
1831 std::size_t range_min_select(std::size_t range_begin,
1832 std::size_t range_end,
1833 std::int64_t target,
1834 std::uint64_t rank) const {
1835 std::int64_t value = 0;
1836 const std::size_t range_end_exclusive = range_end + 1;
1837 const std::size_t first_full_block =
1838 (range_begin + kBlockBits - 1) / kBlockBits;
1839 const std::size_t full_begin =
1840 std::min(range_end_exclusive, first_full_block * kBlockBits);
1841 if (range_begin < full_begin) {
1842 const ScanResult scan = scan_range(range_begin, full_begin);
1843 if (scan.min_value == target) {
1844 if (rank <= scan.min_count) {
1845 return qth_min_in_range(range_begin, full_begin, target, rank);
1846 }
1847 rank -= scan.min_count;
1848 }
1849 value += scan.block_excess;
1850 }
1851
1852 const std::size_t last_full_block_exclusive =
1853 range_end_exclusive / kBlockBits;
1854 const std::size_t middle_begin = full_begin;
1855 const std::size_t middle_end =
1856 std::max(middle_begin, last_full_block_exclusive * kBlockBits);
1857 if (middle_begin < middle_end) {
1858 bool found_node = false;
1859 std::size_t selected = npos;
1860 for_each_cover_node(middle_begin, middle_end, [&](NodeRef node) {
1861 const Summary summary = summary_at(node.level, node.index);
1862 const std::int64_t candidate = value + summary.min_excess;
1863 if (candidate == target) {
1864 if (rank <= summary.min_count) {
1865 selected =
1866 descend_qth_min(node.level, node.index, target - value, rank);
1867 found_node = true;
1868 return false;
1869 }
1870 rank -= summary.min_count;
1871 }
1872 value += summary.block_excess;
1873 return true;
1874 });
1875 if (found_node) {
1876 return selected;
1877 }
1878 }
1879
1880 if (middle_end < range_end_exclusive) {
1881 const ScanResult scan = scan_range(middle_end, range_end_exclusive);
1882 if (value + scan.min_value == target) {
1883 return qth_min_in_range(middle_end, range_end_exclusive, target - value,
1884 rank);
1885 }
1886 }
1887 return npos;
1888 }
1889
1897 template <class Callback>
1898 bool for_each_cover_node(std::size_t begin,
1899 std::size_t end,
1900 Callback&& callback) const {
1901 if (begin >= end || total_levels() == 0 || (begin % kBlockBits) != 0 ||
1902 (end % kBlockBits) != 0) {
1903 return true;
1904 }
1905
1906 std::array<NodeRef, kMaxCoverNodes> right_nodes;
1907 std::size_t right_size = 0;
1908 std::size_t emitted = 0;
1909 std::size_t level = 0;
1910 std::size_t left = begin / kBlockBits;
1911 std::size_t right = end / kBlockBits;
1912
1913 auto emit = [&](NodeRef node) {
1914 if (emitted >= kMaxCoverNodes) {
1915 return true;
1916 }
1917 ++emitted;
1918 return callback(node);
1919 };
1920
1921 auto push_right = [&](NodeRef node) {
1922 if (right_size < right_nodes.size()) {
1923 right_nodes[right_size++] = node;
1924 }
1925 };
1926
1927 while (left < right) {
1928 if (!has_parent_level(level)) {
1929 for (std::size_t index = left; index < right; ++index) {
1930 if (!emit({level, index})) {
1931 return false;
1932 }
1933 }
1934 break;
1935 }
1936
1937 const std::size_t fanout = fanout_to_parent(level);
1938 while (left < right && (left % fanout) != 0) {
1939 if (!emit({level, left})) {
1940 return false;
1941 }
1942 ++left;
1943 }
1944 while (left < right && (right % fanout) != 0) {
1945 --right;
1946 push_right({level, right});
1947 }
1948 left /= fanout;
1949 right /= fanout;
1950 ++level;
1951 }
1952
1953 while (right_size > 0) {
1954 if (!emit(right_nodes[--right_size])) {
1955 return false;
1956 }
1957 }
1958 return true;
1959 }
1960
1971 std::size_t descend_first_extreme(std::size_t level,
1972 std::size_t index,
1973 std::int64_t target,
1974 bool find_min) const {
1975 while (level > 0) {
1976 const std::size_t child_level = level - 1;
1977 const std::size_t fanout = fanout_to_parent(child_level);
1978 const std::size_t child_begin = index * fanout;
1979 const std::size_t child_end =
1980 std::min(level_count(child_level), child_begin + fanout);
1981 std::int64_t prefix = 0;
1982 bool found = false;
1983 for (std::size_t child = child_begin; child < child_end; ++child) {
1984 Summary summary = summary_at(child_level, child);
1985 const std::int64_t candidate =
1986 prefix + (find_min ? summary.min_excess : summary.max_excess);
1987 if (candidate == target) {
1988 index = child;
1989 level = child_level;
1990 target -= prefix;
1991 found = true;
1992 break;
1993 }
1994 prefix += summary.block_excess;
1995 }
1996 if (!found) {
1997 return npos;
1998 }
1999 }
2000 return first_prefix_in_block(index, target);
2001 }
2002
2012 std::size_t first_prefix_in_block(std::size_t block_index,
2013 std::int64_t target) const {
2014 if (full_block_has_words(block_index) && target >= -512 && target <= 512) {
2015 std::uint64_t out[kBlockWords];
2016 excess_positions_512(bits_.data() + block_index * kBlockWords,
2017 static_cast<int>(target), out);
2018 for (std::size_t word = 0; word < kBlockWords; ++word) {
2019 const std::uint64_t mask = out[word];
2020 if (mask != 0) {
2021 return block_index * kBlockBits + word * 64 + std::countr_zero(mask);
2022 }
2023 }
2024 return npos;
2025 }
2026
2027 const std::size_t begin = block_index * kBlockBits;
2028 const std::size_t length = block_size(block_index);
2029 std::int64_t current = 0;
2030 for (std::size_t offset = 0; offset < length; ++offset) {
2031 current += bit(begin + offset) ? 1 : -1;
2032 if (current == target) {
2033 return begin + offset;
2034 }
2035 }
2036 return npos;
2037 }
2038
2049 std::size_t descend_qth_min(std::size_t level,
2050 std::size_t index,
2051 std::int64_t target,
2052 std::uint64_t rank) const {
2053 while (level > 0) {
2054 const std::size_t child_level = level - 1;
2055 const std::size_t fanout = fanout_to_parent(child_level);
2056 const std::size_t child_begin = index * fanout;
2057 const std::size_t child_end =
2058 std::min(level_count(child_level), child_begin + fanout);
2059 std::int64_t prefix = 0;
2060 bool found = false;
2061 for (std::size_t child = child_begin; child < child_end; ++child) {
2062 Summary summary = summary_at(child_level, child);
2063 const std::int64_t candidate = prefix + summary.min_excess;
2064 if (candidate == target) {
2065 if (rank <= summary.min_count) {
2066 index = child;
2067 level = child_level;
2068 target -= prefix;
2069 found = true;
2070 break;
2071 }
2072 rank -= summary.min_count;
2073 }
2074 prefix += summary.block_excess;
2075 }
2076 if (!found) {
2077 return npos;
2078 }
2079 }
2080 return qth_min_in_range(index * kBlockBits,
2081 index * kBlockBits + block_size(index), target,
2082 rank);
2083 }
2084
2095 std::size_t qth_min_in_range(std::size_t begin,
2096 std::size_t end,
2097 std::int64_t target,
2098 std::uint64_t rank) const {
2099 std::int64_t current = 0;
2100 for (std::size_t position = begin; position < end; ++position) {
2101 current += bit(position) ? 1 : -1;
2102 if (current == target && --rank == 0) {
2103 return position;
2104 }
2105 }
2106 return npos;
2107 }
2108
2118 ScanResult scan_range(std::size_t begin, std::size_t end) const {
2119 ScanResult result;
2120 const auto& lut = byte_lut();
2121 while (begin < end && (begin & 7) != 0) {
2122 append_scanned_bit(result, begin);
2123 ++begin;
2124 }
2125 while (begin + 8 <= end) {
2126 const ByteAgg& byte = lut[get_byte(begin)];
2127 const std::int64_t min_candidate = result.block_excess + byte.min_excess;
2128 if (min_candidate < result.min_value) {
2129 result.min_value = min_candidate;
2130 result.min_count = byte.min_count;
2131 result.min_position = begin + byte.pos_first_min;
2132 } else if (min_candidate == result.min_value) {
2133 result.min_count += byte.min_count;
2134 }
2135 const std::int64_t max_candidate = result.block_excess + byte.max_excess;
2136 if (max_candidate > result.max_value) {
2137 result.max_value = max_candidate;
2138 result.max_position = begin + byte.pos_first_max;
2139 }
2140 result.block_excess += byte.block_excess;
2141 begin += 8;
2142 }
2143 while (begin < end) {
2144 append_scanned_bit(result, begin);
2145 ++begin;
2146 }
2147 return result;
2148 }
2149
2155 MinScanResult scan_min_range(std::size_t begin, std::size_t end) const {
2156 MinScanResult result;
2157 const auto& lut = byte_lut();
2158 while (begin < end && (begin & 7) != 0) {
2159 append_scanned_min_bit(result, begin);
2160 ++begin;
2161 }
2162 while (begin + 8 <= end) {
2163 const ByteAgg& byte = lut[get_byte(begin)];
2164 const std::int64_t min_candidate = result.block_excess + byte.min_excess;
2165 if (min_candidate < result.min_value) {
2166 result.min_value = min_candidate;
2167 result.min_position = begin + byte.pos_first_min;
2168 }
2169 result.block_excess += byte.block_excess;
2170 begin += 8;
2171 }
2172 while (begin < end) {
2173 append_scanned_min_bit(result, begin);
2174 ++begin;
2175 }
2176 return result;
2177 }
2178
2186 void append_scanned_bit(ScanResult& result, std::size_t position) const {
2187 result.block_excess += bit(position) ? 1 : -1;
2188 if (result.block_excess < result.min_value) {
2189 result.min_value = result.block_excess;
2190 result.min_count = 1;
2191 result.min_position = position;
2192 } else if (result.block_excess == result.min_value) {
2193 ++result.min_count;
2194 }
2195 if (result.block_excess > result.max_value) {
2196 result.max_value = result.block_excess;
2197 result.max_position = position;
2198 }
2199 }
2200
2204 void append_scanned_min_bit(MinScanResult& result,
2205 std::size_t position) const {
2206 result.block_excess += bit(position) ? 1 : -1;
2207 if (result.block_excess < result.min_value) {
2208 result.min_value = result.block_excess;
2209 result.min_position = position;
2210 }
2211 }
2212
2221 std::uint8_t get_byte(std::size_t bit_position) const {
2222 return static_cast<std::uint8_t>(
2223 (bits_[bit_position >> 6] >> (bit_position & 63)) & 0xffu);
2224 }
2225
2232 std::uint8_t bit(std::size_t position) const {
2233 return static_cast<std::uint8_t>((bits_[position >> 6] >> (position & 63)) &
2234 1ull);
2235 }
2236
2237 std::span<const std::uint64_t> bits_;
2238 std::optional<RankSelectSupport<>> rank_index_;
2239 std::size_t bit_count_ = 0;
2240 std::size_t block_count_ = 0;
2241 Summary top_summary_;
2242 std::vector<std::size_t> level_counts_;
2243 std::vector<std::vector<LowNode>> low_levels_;
2244 std::vector<std::vector<HighNode>> high_levels_;
2245};
2246
2247} // namespace pixie::experimental
SelectSupport
Select directions to index during construction.
Definition support.h:60
CRTP facade for rank/select and range min-max tree operations.
Definition rmm.h:29
RmMBTree(std::span< const std::uint64_t > words, std::size_t bit_count, std::size_t=kBlockBits)
Construct an RmM btree over an external bit-vector span.
Definition btree.h:126
RmMBTree(std::span< const std::uint64_t > words, std::size_t bit_count, RankSelectSupport<>::SelectSupport select_support, std::optional< std::size_t > one_count=std::nullopt, std::size_t=kBlockBits)
Construct with explicit rank/select support options.
Definition btree.h:140
std::size_t memory_usage_bytes() const
Return owned auxiliary memory usage in bytes.
Definition btree.h:342
RangeMinQueryResult range_min_query_result(std::size_t range_begin, std::size_t range_end) const
Return the first minimum position and value in one traversal.
Definition btree.h:310
Definition excess.h:90
Common interface for rank/select and range min-max indexes.
Minimum position and value returned by one range-min traversal.
Definition btree.h:111