1 // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
2 #pragma once
3 
4 #include <algorithm>
5 #include <string>
6 #include <vector>
7 
8 #include "db/dbformat.h"
9 #include "rocksdb/iterator.h"
10 #include "rocksdb/slice.h"
11 #include "table/internal_iterator.h"
12 
13 namespace ROCKSDB_NAMESPACE {
14 
15 // Iterator over a vector of keys/values
16 class VectorIterator : public InternalIterator {
17  public:
VectorIterator(std::vector<std::string> keys,std::vector<std::string> values,const InternalKeyComparator * icmp)18   VectorIterator(std::vector<std::string> keys, std::vector<std::string> values,
19                  const InternalKeyComparator* icmp)
20       : keys_(std::move(keys)),
21         values_(std::move(values)),
22         indexed_cmp_(icmp, &keys_),
23         current_(0) {
24     assert(keys_.size() == values_.size());
25 
26     indices_.reserve(keys_.size());
27     for (size_t i = 0; i < keys_.size(); i++) {
28       indices_.push_back(i);
29     }
30     std::sort(indices_.begin(), indices_.end(), indexed_cmp_);
31   }
32 
Valid()33   virtual bool Valid() const override {
34     return !indices_.empty() && current_ < indices_.size();
35   }
36 
SeekToFirst()37   virtual void SeekToFirst() override { current_ = 0; }
SeekToLast()38   virtual void SeekToLast() override { current_ = indices_.size() - 1; }
39 
Seek(const Slice & target)40   virtual void Seek(const Slice& target) override {
41     current_ = std::lower_bound(indices_.begin(), indices_.end(), target,
42                                 indexed_cmp_) -
43                indices_.begin();
44   }
45 
SeekForPrev(const Slice & target)46   virtual void SeekForPrev(const Slice& target) override {
47     current_ = std::lower_bound(indices_.begin(), indices_.end(), target,
48                                 indexed_cmp_) -
49                indices_.begin();
50     if (!Valid()) {
51       SeekToLast();
52     } else {
53       Prev();
54     }
55   }
56 
Next()57   virtual void Next() override { current_++; }
Prev()58   virtual void Prev() override { current_--; }
59 
key()60   virtual Slice key() const override {
61     return Slice(keys_[indices_[current_]]);
62   }
value()63   virtual Slice value() const override {
64     return Slice(values_[indices_[current_]]);
65   }
66 
status()67   virtual Status status() const override { return Status::OK(); }
68 
IsKeyPinned()69   virtual bool IsKeyPinned() const override { return true; }
IsValuePinned()70   virtual bool IsValuePinned() const override { return true; }
71 
72  private:
73   struct IndexedKeyComparator {
IndexedKeyComparatorIndexedKeyComparator74     IndexedKeyComparator(const InternalKeyComparator* c,
75                          const std::vector<std::string>* ks)
76         : cmp(c), keys(ks) {}
77 
operatorIndexedKeyComparator78     bool operator()(size_t a, size_t b) const {
79       return cmp->Compare((*keys)[a], (*keys)[b]) < 0;
80     }
81 
operatorIndexedKeyComparator82     bool operator()(size_t a, const Slice& b) const {
83       return cmp->Compare((*keys)[a], b) < 0;
84     }
85 
operatorIndexedKeyComparator86     bool operator()(const Slice& a, size_t b) const {
87       return cmp->Compare(a, (*keys)[b]) < 0;
88     }
89 
90     const InternalKeyComparator* cmp;
91     const std::vector<std::string>* keys;
92   };
93 
94   std::vector<std::string> keys_;
95   std::vector<std::string> values_;
96   IndexedKeyComparator indexed_cmp_;
97   std::vector<size_t> indices_;
98   size_t current_;
99 };
100 
101 }  // namespace ROCKSDB_NAMESPACE
102