Program Listing for File rocksdb_set_view.h¶
↰ Return to documentation for file (src/ifcparse/rocksdb_set_view.h)
/********************************************************************************
* *
* This file is part of IfcOpenShell. *
* *
* IfcOpenShell is free software: you can redistribute it and/or modify *
* it under the terms of the Lesser GNU General Public License as published by *
* the Free Software Foundation, either version 3.0 of the License, or *
* (at your option) any later version. *
* *
* IfcOpenShell is distributed in the hope that it will be useful, *
* but WITHOUT ANY WARRANTY; without even the implied warranty of *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the *
* Lesser GNU General Public License for more details. *
* *
* You should have received a copy of the Lesser GNU General Public License *
* along with this program. If not, see <http://www.gnu.org/licenses/>. *
* *
********************************************************************************/
#ifndef ROCKSDB_SET_VIEW_H
#define ROCKSDB_SET_VIEW_H
#ifdef IFOPSH_WITH_ROCKSDB
#include <rocksdb/db.h>
#include <rocksdb/options.h>
#include <rocksdb/write_batch.h>
#endif
#include <memory>
#include <string>
#include <iterator>
#include <cstddef>
#include <cstring>
#include <stdexcept>
template <typename KeyT>
class rocksdb_set_view {
public:
using key_type = KeyT;
using value_type = key_type;
private:
rocksdb::DB* db_;
std::string prefix_;
public:
// Constructor: provide a pointer to an open RocksDB instance and the key-space prefix,
// e.g. "i|"
rocksdb_set_view(rocksdb::DB* db, const std::string& prefix)
: db_(db), prefix_(prefix) {}
// --- Iterator ---
class iterator {
public:
using value_type = key_type;
using difference_type = std::ptrdiff_t;
using iterator_category = std::forward_iterator_tag;
using pointer = const value_type*;
using reference = const value_type&;
private:
rocksdb::DB* db_;
std::string prefix_;
std::unique_ptr<rocksdb::Iterator> it_;
mutable value_type cached_value_;
// Helper: extract the key (i.e. the value) from the current RocksDB key.
value_type extract_current_value() const {
#ifdef IFOPSH_WITH_ROCKSDB
std::string full_key = it_->key().ToString();
std::string remainder = full_key.substr(prefix_.size());
size_t pos = remainder.find('|');
std::string key_str = (pos != std::string::npos) ? remainder.substr(0, pos) : remainder;
return key_from_string<key_type>(key_str);
#else
return cached_value_;
#endif
}
// Validates the current iterator state.
void check_valid() {
#ifdef IFOPSH_WITH_ROCKSDB
if (!it_ || !it_->Valid() || !it_->key().starts_with(prefix_))
it_.reset();
#endif
}
public:
iterator() : db_(nullptr), prefix_(), it_(nullptr) {}
iterator(rocksdb::DB* db, const std::string& prefix,
std::unique_ptr<rocksdb::Iterator> iterator)
: db_(db), prefix_(prefix), it_(std::move(iterator))
{
check_valid();
}
iterator(const iterator& other)
: db_(other.db_), prefix_(other.prefix_)
{
#ifdef IFOPSH_WITH_ROCKSDB
if (other.it_) {
std::string curr = other.it_->key().ToString();
it_.reset(db_->NewIterator(rocksdb::ReadOptions{}));
it_->Seek(curr);
if (!it_->Valid() || it_->key().ToString() != curr)
it_.reset();
}
#endif
}
iterator& operator=(const iterator& other) {
#ifdef IFOPSH_WITH_ROCKSDB
if (this != &other) {
db_ = other.db_;
prefix_ = other.prefix_;
if (other.it_) {
std::string curr = other.it_->key().ToString();
it_.reset(db_->NewIterator(rocksdb::ReadOptions{}));
it_->Seek(curr);
if (!it_->Valid() || it_->key().ToString() != curr)
it_.reset();
} else {
it_.reset();
}
}
#endif
return *this;
}
// Dereference: extract the key part from the RocksDB key.
value_type operator*() const {
return extract_current_value();
}
// Pointer access (via a cached value)
const value_type* operator->() const {
cached_value_ = **this;
return &cached_value_;
}
// Pre-increment: advance the iterator and skip over any duplicate keys.
iterator& operator++() {
#ifdef IFOPSH_WITH_ROCKSDB
if (it_) {
// Record the current key value.
value_type curr = extract_current_value();
do {
it_->Next();
} while (it_ && it_->Valid() && it_->key().starts_with(prefix_) &&
(extract_current_value() == curr));
if (!it_ || !it_->Valid() || !it_->key().starts_with(prefix_))
it_.reset();
}
#endif
return *this;
}
iterator operator++(int) {
iterator tmp(*this);
++(*this);
return tmp;
}
bool operator==(const iterator& other) const {
#ifndef IFOPSH_WITH_ROCKSDB
(void)other;
#endif
#ifdef IFOPSH_WITH_ROCKSDB
if (!it_ && !other.it_)
return true;
if (it_ && other.it_)
return it_->key().ToString() == other.it_->key().ToString();
#endif
return false;
}
bool operator!=(const iterator& other) const {
return !(*this == other);
}
};
// Returns an iterator to the first element in the key-space (or end() if none exist).
iterator begin() const {
#ifdef IFOPSH_WITH_ROCKSDB
auto iter = std::unique_ptr<rocksdb::Iterator>(db_->NewIterator(rocksdb::ReadOptions{}));
iter->Seek(prefix_);
if (iter->Valid() && iter->key().starts_with(prefix_))
return iterator(db_, prefix_, std::move(iter));
#endif
return end();
}
// Returns an iterator representing the end of the key-space.
iterator end() const {
return iterator();
}
// Removes the element: every key under prefix + key + "|". The exclusive
// upper bound is the same prefix with its separator incremented. Returns
// 1 if the element existed, 0 otherwise.
size_t erase(const key_type& key) {
#ifdef IFOPSH_WITH_ROCKSDB
if (find(key) == end()) {
return 0;
}
const std::string lower_bound = prefix_ + key_to_string(key) + "|";
const std::string upper_bound = prefix_ + key_to_string(key) + std::string(1, '|' + 1);
rocksdb::WriteBatch batch;
batch.DeleteRange(lower_bound, upper_bound);
db_->Write(rocksdb::WriteOptions{}, &batch);
return 1;
#else
static_cast<void>(key);
return 0;
#endif
}
// Read-only find: returns an iterator to the element with the given key if it exists.
iterator find(const key_type& key) const {
#ifdef IFOPSH_WITH_ROCKSDB
std::string key_str = key_to_string(key);
// Construct the search key: prefix + key_str + separator.
std::string start_key = prefix_ + key_str + "|";
auto iter = std::unique_ptr<rocksdb::Iterator>(db_->NewIterator(rocksdb::ReadOptions{}));
iter->Seek(start_key);
if (iter->Valid() && iter->key().starts_with(prefix_)) {
std::string full_key = iter->key().ToString();
std::string remainder = full_key.substr(prefix_.size());
size_t pos = remainder.find('|');
std::string found_key_str = (pos != std::string::npos) ? remainder.substr(0, pos) : remainder;
if (key_from_string<key_type>(found_key_str) == key)
return iterator(db_, prefix_, std::move(iter));
}
#endif
return end();
}
};
#endif