#pragma once
#include <array>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>
#include "common/type_utils.h"
#include "common/types/int128_t.h"
#include "common/types/string_t.h"
#include "common/types/uint128_t.h"
#include "storage/index/index.h"
namespace lbug {
namespace common {
struct BufferReader;
}
namespace storage {
class ArtKey {
public:
ArtKey() = default;
explicit ArtKey(std::vector<uint8_t> bytes) : bytes{std::move(bytes)} {}
bool empty() const { return bytes.empty(); }
const std::vector<uint8_t>& getBytes() const { return bytes; }
static ArtKey encode(common::ValueVector* vector, uint64_t vectorPos);
private:
std::vector<uint8_t> bytes;
};
struct ArtPrimaryKeyIndexStorageInfo final : IndexStorageInfo {
std::vector<std::pair<std::vector<uint8_t>, common::offset_t>> entries;
ArtPrimaryKeyIndexStorageInfo() = default;
explicit ArtPrimaryKeyIndexStorageInfo(
std::vector<std::pair<std::vector<uint8_t>, common::offset_t>> entries)
: entries{std::move(entries)} {}
DELETE_COPY_DEFAULT_MOVE(ArtPrimaryKeyIndexStorageInfo);
std::shared_ptr<common::BufferWriter> serialize() const override;
static std::unique_ptr<IndexStorageInfo> deserialize(
std::unique_ptr<common::BufferReader> reader);
};
class ArtPrimaryKeyIndex final : public Index {
public:
struct InsertState final : Index::InsertState {
visible_func isVisible;
explicit InsertState(visible_func isVisible) : isVisible{std::move(isVisible)} {}
};
explicit ArtPrimaryKeyIndex(IndexInfo indexInfo, std::unique_ptr<IndexStorageInfo> storageInfo);
~ArtPrimaryKeyIndex() override;
static std::unique_ptr<ArtPrimaryKeyIndex> createNewIndex(IndexInfo indexInfo);
std::unique_ptr<Index::InsertState> initInsertState(main::ClientContext*,
visible_func isVisible) override;
bool needCommitInsert() const override { return true; }
void commitInsert(transaction::Transaction* transaction,
const common::ValueVector& nodeIDVector,
const std::vector<common::ValueVector*>& indexVectors,
Index::InsertState& insertState) override;
std::unique_ptr<DeleteState> initDeleteState(const transaction::Transaction*, MemoryManager*,
visible_func) override {
return std::make_unique<DeleteState>();
}
void delete_(transaction::Transaction*, const common::ValueVector&, DeleteState&) override {
}
bool lookupPrimaryKey(const transaction::Transaction* transaction,
common::ValueVector* keyVector, uint64_t vectorPos, common::offset_t& result,
visible_func isVisible) override;
bool scanPrimaryKeyRange(common::ValueVector* lowerBoundVector, uint64_t lowerBoundPos,
bool lowerInclusive, common::ValueVector* upperBoundVector, uint64_t upperBoundPos,
bool upperInclusive, common::idx_t maxResults, std::vector<common::offset_t>& results,
visible_func isVisible) override;
void discardPrimaryKey(common::ValueVector* keyVector) override;
void checkpoint(main::ClientContext*, PageAllocator&) override;
static LBUG_API std::unique_ptr<Index> load(main::ClientContext* context,
StorageManager* storageManager, IndexInfo indexInfo, std::span<uint8_t> storageInfoBuffer);
static IndexType getIndexType() {
static const IndexType ART_INDEX_TYPE{"ART", IndexConstraintType::PRIMARY,
IndexDefinitionType::BUILTIN, load};
return ART_INDEX_TYPE;
}
private:
struct Node {
static constexpr uint8_t EMPTY_MARKER = UINT8_MAX;
enum class Kind : uint8_t { NODE4 = 0, NODE16 = 1, NODE48 = 2, NODE256 = 3 };
struct SmallChildren {
std::array<uint8_t, 16> keys{};
std::array<Node*, 16> children{};
};
struct Node48Children {
std::array<uint8_t, 256> childIndex{};
std::array<Node*, 48> children{};
};
struct Node256Children {
std::array<Node*, 256> children{};
};
std::optional<common::offset_t> offset;
Kind kind = Kind::NODE4;
uint16_t count = 0;
union {
SmallChildren small;
Node48Children node48;
Node256Children node256;
};
Node();
Node* getChild(uint8_t byte) const;
Node* getOrInsertChild(ArtPrimaryKeyIndex& index, uint8_t byte);
void removeChild(uint8_t byte);
bool empty() const { return !offset.has_value() && count == 0; }
};
static constexpr uint64_t NODE_BLOCK_CAPACITY = 16 * 1024;
struct NodeBlock {
Node* nodes = nullptr;
uint64_t used = 0;
NodeBlock();
~NodeBlock();
NodeBlock(const NodeBlock&) = delete;
NodeBlock& operator=(const NodeBlock&) = delete;
NodeBlock(NodeBlock&& other) noexcept;
NodeBlock& operator=(NodeBlock&& other) noexcept;
};
bool insertInternal(const ArtKey& key, common::offset_t offset, visible_func isVisible);
bool lookup(const ArtKey& key, common::offset_t& result, visible_func isVisible) const;
bool eraseInternal(Node& node, const std::vector<uint8_t>& key, uint64_t depth);
void erase(const ArtKey& key);
Node* allocateNode();
void recordKindChange(Node& node, Node::Kind newKind);
void collectRange(const Node& node, std::vector<uint8_t>& key, const ArtKey* lowerBound,
bool lowerInclusive, const ArtKey* upperBound, bool upperInclusive,
common::idx_t maxResults, std::vector<common::offset_t>& results,
visible_func isVisible) const;
void clear();
void collectEntries(const Node& node, std::vector<uint8_t>& key,
std::vector<std::pair<std::vector<uint8_t>, common::offset_t>>& entries) const;
void loadEntries(const ArtPrimaryKeyIndexStorageInfo& storageInfo);
private:
Node root;
std::vector<NodeBlock> nodeBlocks;
uint64_t numAllocatedNodes = 1;
std::array<uint64_t, 4> numNodesByKind{1, 0, 0, 0};
mutable std::mutex mutex;
};
} }