Program Listing for File rocksdb_map_adapter.h

↰ Return to documentation for file (src/ifcparse/rocksdb_map_adapter.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_MAP_ADAPTER_H
#define ROCKSDB_MAP_ADAPTER_H

#ifdef IFOPSH_WITH_ROCKSDB
#include <rocksdb/db.h>
#include <rocksdb/options.h>
#endif

#include <memory>
#include <string>
#include <utility>
#include <iterator>
#include <cstddef>
#include <cstdint>
#include <cstring>

template <typename T>
struct is_std_tuple : std::false_type {};

template <typename... Ts>
struct is_std_tuple<std::tuple<Ts...>> : std::true_type {};

// Serialization and deserialization primitives
template <typename T>
struct DefaultCodec;

// @todo specialize for integral types in one go

// Specialization for size_t.
template <>
struct DefaultCodec<size_t> {
    std::string encode(const size_t& value) const {
        std::string encoded(sizeof(value), 0);
        memcpy(encoded.data(), &value, sizeof(value));
        return encoded;
    }
    size_t decode(const std::string& encoded) const {
        size_t value = 0;
        // @todo take min of sizeof(value), len(encoded)
        // @todo unify all serialization primitives
        memcpy(&value, encoded.data(), sizeof(value));
        return value;
    }
};

// Specialization for uint32_t.
template <>
struct DefaultCodec<uint32_t> {
    std::string encode(const uint32_t& value) const {
        std::string encoded(sizeof(value), 0);
        memcpy(encoded.data(), &value, sizeof(value));
        return encoded;
    }
    uint32_t decode(const std::string& encoded) const {
        uint32_t value = 0;
        // @todo take min of sizeof(value), len(encoded)
        // @todo unify all serialization primitives
        memcpy(&value, encoded.data(), sizeof(value));
        return value;
    }
};

// Specialization for std::vector<int>.
template <>
struct DefaultCodec<std::vector<uint32_t>> {
    std::string encode(const std::vector<uint32_t>& values) const {
        std::string encoded(sizeof(uint32_t) * values.size(), 0);
        memcpy(encoded.data(), values.data(), encoded.size());
        return encoded;
    }
    std::vector<uint32_t> decode(const std::string& encoded) const {
        std::vector<uint32_t> values(encoded.size() / sizeof(uint32_t), 0);
        memcpy(values.data(), encoded.data(), encoded.size());
        return values;
    }
};

// Specialization for std::string (identity)
template <>
struct DefaultCodec<std::string> {
    std::string encode(const std::string& value) const {
        return value;
    }
    std::string decode(const std::string& encoded) const {
        return encoded;
    }
};

// Numeric key segments are fixed-width lowercase hex, so keys sort as
// numbers: an instance's records are laid out in id order and the largest
// id under a prefix is its last key.
inline std::string hex_key(uint64_t value) {
    static constexpr char digits[] = "0123456789abcdef";
    std::string s(16, '0');
    for (int i = 15; i >= 0; --i) {
        s[(size_t)i] = digits[value & 0xf];
        value >>= 4;
    }
    return s;
}

inline uint64_t parse_hex_key(const std::string& s) {
    return std::stoull(s, nullptr, 16);
}

template <typename KeyT>
std::string key_to_string(const KeyT& key) {
    if constexpr (std::is_same_v<KeyT, std::string>) {
        return key;
    } else {
        static_assert(std::is_integral_v<KeyT>, "key_to_string expects a string or an integral key");
        return hex_key((uint64_t)key);
    }
}

// The keys the file storage and the serializer agree on:
//   i|<id>|<attribute>   t|<identity>|<attribute>   h|<name>|<attribute>
//   i|<id>|_             t|<identity>|_             type record
//   v|<referenced id>|<entity index>|<attribute>    inverse record
//   t|<declaration index>                           instances of a type
namespace rocksdb_key {
    // The exclusive end of everything under prefix.
    inline std::string upper_bound(std::string prefix) {
        prefix.back() = (char)(prefix.back() + 1);
        return prefix;
    }
    inline std::string instance(bool is_entity, uint64_t id) {
        return (is_entity ? "i|" : "t|") + key_to_string(id) + "|";
    }
    inline std::string attribute(bool is_entity, uint64_t id, uint64_t index) {
        return instance(is_entity, id) + key_to_string(index);
    }
    inline std::string header_attribute(const std::string& name, uint64_t index) {
        return "h|" + name + "|" + key_to_string(index);
    }
    inline std::string type_record(bool is_entity, uint64_t id) {
        return instance(is_entity, id) + "_";
    }
    inline std::string inverse_prefix(uint64_t referenced_id) {
        return "v|" + key_to_string(referenced_id) + "|";
    }
    inline std::string inverse(uint64_t referenced_id, uint64_t entity_index, uint64_t attribute_index) {
        return inverse_prefix(referenced_id) + key_to_string(entity_index) + "|" + key_to_string(attribute_index);
    }
    inline std::string type_list(uint64_t declaration_index) {
        return "t|" + key_to_string(declaration_index);
    }
}

// Convert from a string to a key. For non-string types, we assume numeric keys.
template <typename KeyT, typename std::enable_if<!is_std_tuple<KeyT>::value, int>::type = 0>
KeyT key_from_string(const std::string& key_string) {
    // @todo tuples
    if constexpr (std::is_same_v<KeyT, std::string>) {
        return key_string;
    } else if constexpr (std::is_integral_v<KeyT>) {
        return static_cast<KeyT>(parse_hex_key(key_string));
    } else {
        static_assert(sizeof(KeyT) == 0, "key_from_string not implemented for this type");
    }
}

template<typename Tuple, std::size_t... Is>
std::string tuple_to_string_impl(const Tuple& tuple_value, std::index_sequence<Is...> indices) {
    static_cast<void>(indices);
    std::ostringstream oss;
    // Unpack the tuple; add a pipe before each element except the first.
    ((oss << (Is == 0 ? "" : "|") << key_to_string(std::get<Is>(tuple_value))), ...);
    return oss.str();
}

template<typename... Ts>
std::string key_to_string(const std::tuple<Ts...>& key) {
    return tuple_to_string_impl(key, std::index_sequence_for<Ts...>{});
}

// Helper: Convert a string token to the desired numeric type.
template<typename T>
T convert_string(const std::string& token) {
    if constexpr (std::is_integral_v<T>) {
        return static_cast<T>(parse_hex_key(token));
    } else if constexpr (std::is_floating_point_v<T>) {
        return static_cast<T>(std::stod(token));
    } else {
        static_assert(sizeof(T) == 0, "convert_string not implemented for this type");
    }
}

// Helper: Build a tuple from a vector of string tokens.
template <typename TupleT, std::size_t... Is>
TupleT tuple_from_string_impl(const std::vector<std::string>& tokens, std::index_sequence<Is...> indices) {
    static_cast<void>(indices);
    return std::make_tuple(convert_string<std::tuple_element_t<Is, TupleT>>(tokens[Is])...);
}

template <typename TupleT, typename std::enable_if<is_std_tuple<TupleT>::value, int>::type = 0>
TupleT key_from_string(const std::string& key_string) {
    std::vector<std::string> tokens;
    std::istringstream iss(key_string);
    std::string token;
    while (std::getline(iss, token, '|')) {
        tokens.push_back(token);
    }
    if (tokens.size() != std::tuple_size<TupleT>::value) {
        throw std::runtime_error("Invalid tuple format");
    }
    return tuple_from_string_impl<TupleT>(tokens, std::make_index_sequence<std::tuple_size<TupleT>::value>{});
}

// rocksdb_map_adapter: a std::map-like interface on a RocksDB keyspace with a given prefix.
// The mapped_type is templated and encoded/decoded via Codec.
template <typename KeyT, typename MappedT, typename Codec = DefaultCodec<MappedT>>
class rocksdb_map_adapter {
public:
    using key_type = KeyT;
    using mapped_type = MappedT;
    using value_type = std::pair<key_type, mapped_type>;

private:
    rocksdb::DB* db_;
    std::string prefix_;
    Codec codec_;

public:
    rocksdb_map_adapter(rocksdb::DB* db, const std::string& prefix)
        : db_(db), prefix_(prefix), codec_(Codec{}) {}

    class iterator {
    public:
        using value_type = std::pair<key_type, mapped_type>;
        using difference_type = std::ptrdiff_t;
        using iterator_category = std::forward_iterator_tag;
        using pointer = value_type*;
        using reference = value_type&;

    private:
        rocksdb::DB* db_;
        std::string prefix_;
        Codec codec_;
        // When it_ is nullptr, this iterator is at end.
        std::unique_ptr<rocksdb::Iterator> it_;
        mutable value_type cached_value_;

        void check_valid() {
#ifdef IFOPSH_WITH_ROCKSDB
            if (!it_ || !it_->Valid() || !it_->key().starts_with(prefix_)) {
                it_.reset();
            }
#endif
        }

    public:
        iterator() : db_(nullptr), prefix_(), codec_(Codec{}), it_(nullptr) {}

        iterator(rocksdb::DB* db, const std::string& prefix,
            std::unique_ptr<rocksdb::Iterator> iterator, Codec codec = Codec{})
            : db_(db), prefix_(prefix), codec_(codec), it_(std::move(iterator))
        {
            check_valid();
        }

        iterator(const iterator& other)
            : db_(other.db_), prefix_(other.prefix_), codec_(other.codec_)
        {
#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_;
                codec_ = other.codec_;
                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;
        }

        value_type operator*() const {
#ifdef IFOPSH_WITH_ROCKSDB
            std::string full_key = it_->key().ToString();
            std::string key_without_prefix = full_key.substr(prefix_.size());
            std::string value_str = it_->value().ToString();
            return { key_from_string<key_type>(key_without_prefix), codec_.decode(value_str) };
#else
            return cached_value_;
#endif
        }

        // operator-> uses a mutable cache to return a pointer to the current value.
        value_type* operator->() const {
            cached_value_ = **this;
            return &cached_value_;
        }

        iterator& operator++() {
#ifdef IFOPSH_WITH_ROCKSDB
            if (it_) {
                it_->Next();
                check_valid();
            }
#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);
        }
    };

    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), codec_);
        }
#endif
        return end();
    }

    iterator end() const {
        return iterator();
    }

        iterator find(const key_type& key) const {
#ifndef IFOPSH_WITH_ROCKSDB
            (void)key;
#endif
#ifdef IFOPSH_WITH_ROCKSDB
        std::string key_str = key_to_string(key);
        std::string full_key = prefix_ + key_str;
        auto iter = std::unique_ptr<rocksdb::Iterator>(db_->NewIterator(rocksdb::ReadOptions{}));
        iter->Seek(full_key);
        if (iter->Valid() && iter->key().ToString() == full_key)
            return iterator(db_, prefix_, std::move(iter), codec_);
#endif
        return end();
    }

        size_t erase(const key_type& key) {
#ifndef IFOPSH_WITH_ROCKSDB
            (void)key;
#endif
#ifdef IFOPSH_WITH_ROCKSDB
        std::string key_str = key_to_string(key);
        std::string full_key = prefix_ + key_str;
        rocksdb::Status s = db_->Delete(rocksdb::WriteOptions{}, full_key);
        return s.ok() ? 1 : 0;
#else
        return 0;
#endif
    }

    std::pair<iterator, bool> insert(const value_type& value) {
#ifdef IFOPSH_WITH_ROCKSDB
        std::string key_str = key_to_string(value.first);
        std::string full_key = prefix_ + key_str;
        std::string existing;
        rocksdb::Status s = db_->Get(rocksdb::ReadOptions{}, full_key, &existing);
        if (s.ok()) {
            // Key already exists.
            return { find(value.first), false };
        }
        std::string encoded = codec_.encode(value.second);
        s = db_->Put(rocksdb::WriteOptions{}, full_key, encoded);
        if (!s.ok()) {
            return { end(), false };
        }
#endif
        return { find(value.first), true };
    }
};

#endif