store.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
 
#ifndef LIBZARR_STORE_HPP
#define LIBZARR_STORE_HPP
 
#include <cstddef>
#include <cstdint>
#include <map>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include "libzarr/types.hpp"
 
 
namespace zarr {
 
struct ByteRange {
  enum class Kind : std::uint8_t {
    full,    
    slice,   
    suffix,  
  };
 
  Kind kind = Kind::full;
  std::uint64_t offset = 0;
  std::uint64_t length = 0;
 
  [[nodiscard]] static constexpr ByteRange full() { return ByteRange{}; }
 
  [[nodiscard]] static constexpr ByteRange slice(std::uint64_t offset, std::uint64_t length) {
    return ByteRange{Kind::slice, offset, length};
  }
 
  [[nodiscard]] static constexpr ByteRange suffix(std::uint64_t length) {
    return ByteRange{Kind::suffix, 0, length};
  }
};
 
struct ReadRequest {
  std::string_view key;
  ByteRange range = ByteRange::full();
};
 
struct DirListing {
  std::vector<std::string> keys;
  std::vector<std::string> prefixes;
};
 
class Store {
 public:
  Store() = default;
  Store(const Store&) = delete;
  Store& operator=(const Store&) = delete;
  Store(Store&&) = delete;
  Store& operator=(Store&&) = delete;
  virtual ~Store() = default;
 
  [[nodiscard]] virtual std::optional<Bytes> read(std::string_view key) = 0;
 
  [[nodiscard]] virtual std::optional<Bytes> read_range(std::string_view key, ByteRange range);
 
  [[nodiscard]] virtual std::optional<std::uint64_t> size(std::string_view key) {
    const auto value = read(key);
    if (!value) {
      return std::nullopt;
    }
    return value->size();
  }
 
  [[nodiscard]] virtual std::vector<std::optional<Bytes>> read_many(
      const std::vector<ReadRequest>& requests) {
    std::vector<std::optional<Bytes>> out;
    out.reserve(requests.size());
    for (const ReadRequest& req : requests) {
      out.push_back(read_range(req.key, req.range));
    }
    return out;
  }
 
  virtual void write(std::string_view key, Bytes value) = 0;
 
  [[nodiscard]] virtual bool exists(std::string_view key) = 0;
 
  virtual void erase(std::string_view key) = 0;
 
  virtual void flush() {}
 
  [[nodiscard]] virtual std::vector<std::string> list_prefix(std::string_view prefix) = 0;
 
  [[nodiscard]] virtual DirListing list_dir(std::string_view prefix) = 0;
};
 
class MemoryStore final : public Store {
 public:
  [[nodiscard]] std::optional<Bytes> read(std::string_view key) override {
    const auto it = map_.find(key);
    if (it == map_.end()) {
      return std::nullopt;
    }
    return it->second;
  }
 
  void write(std::string_view key, Bytes value) override {
    map_.insert_or_assign(std::string(key), std::move(value));
  }
 
  [[nodiscard]] std::optional<std::uint64_t> size(std::string_view key) override {
    const auto it = map_.find(key);
    if (it == map_.end()) {
      return std::nullopt;
    }
    return it->second.size();
  }
 
  [[nodiscard]] bool exists(std::string_view key) override { return map_.find(key) != map_.end(); }
 
  void erase(std::string_view key) override {
    const auto it = map_.find(key);
    if (it != map_.end()) {
      map_.erase(it);
    }
  }
 
  [[nodiscard]] std::vector<std::string> list_prefix(std::string_view prefix) override {
    check_prefix(prefix);
    std::vector<std::string> out;
    for (auto it = map_.lower_bound(prefix); it != map_.end() && starts_with(it->first, prefix);
         ++it) {
      out.push_back(it->first);
    }
    return out;
  }
 
  [[nodiscard]] DirListing list_dir(std::string_view prefix) override {
    check_prefix(prefix);
    DirListing out;
    for (auto it = map_.lower_bound(prefix); it != map_.end() && starts_with(it->first, prefix);
         ++it) {
      const auto rest = std::string_view(it->first).substr(prefix.size());
      const auto slash = rest.find('/');
      if (slash == std::string_view::npos) {
        out.keys.emplace_back(rest);
      } else {
        // Keys under one child prefix are contiguous in a sorted map, so
        // comparing against the last emitted prefix deduplicates.
        const auto child = rest.substr(0, slash);
        if (out.prefixes.empty() || out.prefixes.back() != child) {
          out.prefixes.emplace_back(child);
        }
      }
    }
    return out;
  }
 
  [[nodiscard]] std::size_t key_count() const { return map_.size(); }
 
 private:
  static bool starts_with(std::string_view text, std::string_view prefix) {
    return text.size() >= prefix.size() && text.compare(0, prefix.size(), prefix) == 0;
  }
 
  static void check_prefix(std::string_view prefix) {
    if (!prefix.empty() && prefix.back() != '/') {
      throw error("store prefix must be empty or end with '/', got '" + std::string(prefix) + "'");
    }
  }
 
  std::map<std::string, Bytes, std::less<>> map_;
};
 
inline std::optional<Bytes> Store::read_range(std::string_view key, ByteRange range) {
  auto value = read(key);
  if (!value) {
    return std::nullopt;
  }
  const std::uint64_t size = value->size();
  std::uint64_t begin = 0;
  std::uint64_t count = 0;
  switch (range.kind) {
    case ByteRange::Kind::full:
      return value;
    case ByteRange::Kind::slice:
      if (range.length > size || range.offset > size - range.length) {
        throw error("read_range: slice at offset " + std::to_string(range.offset) + " of length " +
                    std::to_string(range.length) + " out of bounds for \"" + std::string(key) +
                    "\" (" + std::to_string(size) + " bytes)");
      }
      begin = range.offset;
      count = range.length;
      break;
    case ByteRange::Kind::suffix:
      if (range.length > size) {
        throw error("read_range: suffix of length " + std::to_string(range.length) +
                    " out of bounds for \"" + std::string(key) + "\" (" + std::to_string(size) +
                    " bytes)");
      }
      begin = size - range.length;
      count = range.length;
      break;
  }
  const auto first = value->begin() + static_cast<std::ptrdiff_t>(begin);
  return Bytes(first, first + static_cast<std::ptrdiff_t>(count));
}
 
}  // namespace zarr
 
#endif  // LIBZARR_STORE_HPP