BKE_bake_items_serialize.hh

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

 
#pragma once
 
#include "BLI_fileops.hh"
#include "BLI_function_ref.hh"
#include "BLI_mutex.hh"
#include "BLI_serialize.hh"
 
#include "BKE_bake_items.hh"
 
namespace blender::bke::bake {

struct BlobSlice {
  std::string name;
  IndexRange range;
 
  std::shared_ptr<io::serialize::DictionaryValue> serialize() const;
  static std::optional<BlobSlice> deserialize(const io::serialize::DictionaryValue &io_slice);
};

class BlobReader {
 public:
  virtual ~BlobReader() = default;

  [[nodiscard]] virtual bool read(const BlobSlice &slice, void *r_data) const = 0;

  [[nodiscard]] virtual bool read_as_stream(const BlobSlice &slice,
                                            FunctionRef<bool(std::istream &)> fn) const;
};

class BlobWriter {
 protected:
  int64_t total_written_size_ = 0;
 
 public:
  virtual ~BlobWriter() = default;

  virtual BlobSlice write(const void *data, int64_t size) = 0;

  virtual BlobSlice write_as_stream(StringRef file_extension,
                                    FunctionRef<void(std::ostream &)> fn);
 
  int64_t written_size() const
  {
    return total_written_size_;
  }
};

class BlobWriteSharing : NonCopyable, NonMovable {
 private:
  struct StoredByRuntimeValue {
    int64_t sharing_info_version;
    std::shared_ptr<io::serialize::DictionaryValue> io_data;
  };

  Map<const ImplicitSharingInfo *, StoredByRuntimeValue> stored_by_runtime_;

  Map<uint64_t, BlobSlice> slice_by_content_hash_;
 
 public:
  ~BlobWriteSharing();

  [[nodiscard]] std::shared_ptr<io::serialize::DictionaryValue> write_implicitly_shared(
      const ImplicitSharingInfo *sharing_info,
      FunctionRef<std::shared_ptr<io::serialize::DictionaryValue>()> write_fn);

  [[nodiscard]] std::shared_ptr<io::serialize::DictionaryValue> write_deduplicated(
      BlobWriter &writer, const void *data, int64_t size_in_bytes);
};

class BlobReadSharing : NonCopyable, NonMovable {
 private:
  mutable Mutex mutex_;
  mutable Map<std::string, ImplicitSharingInfoAndData> runtime_by_stored_;
 
 public:
  ~BlobReadSharing();

  [[nodiscard]] std::optional<ImplicitSharingInfoAndData> read_shared(
      const io::serialize::DictionaryValue &io_data,
      FunctionRef<std::optional<ImplicitSharingInfoAndData>()> read_fn) const;
};

class DiskBlobReader : public BlobReader {
 private:
  const std::string blobs_dir_;
  mutable Mutex mutex_;
  mutable Map<std::string, std::unique_ptr<fstream>> open_input_streams_;
 
 public:
  DiskBlobReader(std::string blobs_dir);
  [[nodiscard]] bool read(const BlobSlice &slice, void *r_data) const override;
};

class DiskBlobWriter : public BlobWriter {
 private:
  std::string blob_dir_;
  std::string base_name_;
  std::string blob_name_;
  std::fstream blob_stream_;
  int64_t current_offset_ = 0;
  int independent_file_count_ = 0;
 
 public:
  DiskBlobWriter(std::string blob_dir, std::string base_name);
 
  BlobSlice write(const void *data, int64_t size) override;
 
  BlobSlice write_as_stream(StringRef file_extension,
                            FunctionRef<void(std::ostream &)> fn) override;
};

class MemoryBlobWriter : public BlobWriter {
 public:
  struct OutputStream {
    std::unique_ptr<std::ostringstream> stream;
    int64_t offset = 0;
  };
 
 private:
  std::string base_name_;
  std::string blob_name_;
  Map<std::string, OutputStream> stream_by_name_;
  int independent_file_count_ = 0;
 
 public:
  MemoryBlobWriter(std::string base_name);
 
  BlobSlice write(const void *data, int64_t size) override;
 
  BlobSlice write_as_stream(StringRef file_extension,
                            FunctionRef<void(std::ostream &)> fn) override;
 
  const Map<std::string, OutputStream> &get_stream_by_name() const
  {
    return stream_by_name_;
  }
};

class MemoryBlobReader : public BlobReader {
 private:
  Map<std::string, Span<std::byte>> blob_by_name_;
 
 public:
  void add(StringRef name, Span<std::byte> blob);
 
  [[nodiscard]] bool read(const BlobSlice &slice, void *r_data) const override;
};
 
void serialize_bake(const BakeState &bake_state,
                    BlobWriter &blob_writer,
                    BlobWriteSharing &blob_sharing,
                    std::ostream &r_stream);
 
std::optional<BakeState> deserialize_bake(std::istream &stream,
                                          const BlobReader &blob_reader,
                                          const BlobReadSharing &blob_sharing);
 
}  // namespace blender::bke::bake