DataSource.h

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <memory>
#include <type_traits>
#include <vector>
 
#include <vrs/DataSlice.h>
#include <vrs/VrsExport.h>
 
namespace vrs {
 
using std::is_pointer;
using std::vector;
 
class DataLayout;
class WriteFileHandler;
 
struct VRS_API DataLayoutChunk {
  DataLayoutChunk() = default;
  explicit DataLayoutChunk(DataLayout& dataLayout);
 
  size_t size() const {
    return layoutFixedSize_ + layoutVariableSize_;
  }
  void fillAndAdvanceBuffer(uint8_t*& buffer) const;
 
  DataLayout* const dataLayout_{};
  const size_t layoutFixedSize_{};
  const size_t layoutVariableSize_{};
};
 
class VRS_API DataSourceChunk {
 public:
  DataSourceChunk(const DataSourceChunk& other) = default;
  DataSourceChunk(DataSourceChunk&& other) noexcept = default;
 
  DataSourceChunk() = default;
  DataSourceChunk(const void* _data, size_t _size);
  template <typename T>
  /* implicit */ DataSourceChunk(const vector<T>& vectorT)
      : data_{vectorT.data()}, size_{sizeof(T) * vectorT.size()} {}
  virtual ~DataSourceChunk() = default;
 
  DataSourceChunk& operator=(const DataSourceChunk& other) = default;
  DataSourceChunk& operator=(DataSourceChunk&& other) noexcept = default;
 
  template <
      typename T,
      // T must be trivially copyable.
      std::enable_if_t<std::is_trivially_copyable<T>::value, int> = 0,
      // T may not be a pointer.
      typename = typename std::enable_if<!is_pointer<T>::value, T>::type>
  /* implicit */ DataSourceChunk(const T& object) : data_{&object}, size_{sizeof(T)} {}
 
  virtual void fillAndAdvanceBuffer(uint8_t*& buffer) const;
 
  const void* data() const {
    return data_;
  }
 
  size_t size() const {
    return size_;
  }
 
 private:
  const void* data_{};
  size_t size_{0};
};
 
class VRS_API NonContiguousChunk final : public vrs::DataSourceChunk {
 public:
  NonContiguousChunk(const void* data, size_t blockSize, size_t numBlocks, size_t strideInBytes);
 
  void fillAndAdvanceBuffer(uint8_t*& buffer) const final;
 
 private:
  const size_t blockSize_;
  const size_t numBlocks_;
  const size_t strideInBytes_;
};
 
class VRS_API DataSource {
 public:
  explicit DataSource(size_t size = 0) : size_{size} {}
  virtual ~DataSource() = default;
  explicit DataSource(const DataSourceChunk& chunk) : chunk1_{chunk}, size_{chunk.size()} {}
  explicit DataSource(const DataSourceChunk& firstChunk, const DataSourceChunk& secondChunk)
      : chunk1_{firstChunk}, chunk2_{secondChunk}, size_{firstChunk.size() + secondChunk.size()} {}
  explicit DataSource(
      const DataSourceChunk& firstChunk,
      const DataSourceChunk& secondChunk,
      const DataSourceChunk& thirdChunk)
      : chunk1_{firstChunk}, chunk2_{secondChunk}, chunk3_{thirdChunk}, size_{getChunksSize()} {}
  explicit DataSource(DataLayout& dataLayout)
      : dataLayout1_{dataLayout}, size_{dataLayout1_.size()} {}
  explicit DataSource(DataLayout& dataLayout, const DataSourceChunk& chunk)
      : dataLayout1_{dataLayout}, chunk1_{chunk}, size_{dataLayout1_.size() + chunk1_.size()} {}
  explicit DataSource(
      DataLayout& dataLayout,
      const DataSourceChunk& firstChunk,
      const DataSourceChunk& secondChunk);
  explicit DataSource(
      DataLayout& dataLayout,
      const DataSourceChunk& firstChunk,
      const DataSourceChunk& secondChunk,
      const DataSourceChunk& thirdChunk);
  explicit DataSource(DataLayout& dataLayout1, DataLayout& dataLayout2);
  explicit DataSource(
      DataLayout& dataLayout1,
      DataLayout& dataLayout2,
      const DataSourceChunk& chunk);
  explicit DataSource(
      DataLayout& dataLayout1,
      DataLayout& dataLayout2,
      const DataSourceChunk& firstChunk,
      const DataSourceChunk& secondChunk);
  explicit DataSource(
      DataLayout& dataLayout1,
      DataLayout& dataLayout2,
      const DataSourceChunk& firstChunk,
      const DataSourceChunk& secondChunk,
      const DataSourceChunk& thirdChunk);
 
  DataSource(const DataSource&) = delete;
  DataSource(DataSource&&) = delete;
 
  DataSource& operator=(const DataSource&) = delete;
  DataSource& operator=(DataSource&&) = delete;
 
  size_t getDataSize() const {
    return size_;
  }
 
  virtual void copyTo(uint8_t* buffer) const;
 
 protected:
  static DataSourceChunk kEmptyDataSourceChunk;
 
  size_t getChunksSize();
 
  const DataLayoutChunk dataLayout1_;
  const DataLayoutChunk dataLayout2_;
  const DataSourceChunk& chunk1_{kEmptyDataSourceChunk};
  const DataSourceChunk& chunk2_{kEmptyDataSourceChunk};
  const DataSourceChunk& chunk3_{kEmptyDataSourceChunk};
  const size_t size_;
};
 
class VRS_API DirectWriteRecordData {
 public:
  virtual ~DirectWriteRecordData() = default;
 
  virtual size_t getDataSize() const = 0;
 
  virtual int write(WriteFileHandler& file) const = 0;
};
 
class VRS_API BorrowedData {
 public:
  BorrowedData(const void* data, size_t size, std::shared_ptr<void> owner)
      : segments_{{data, size}}, owner_{std::move(owner)}, size_{size} {}
 
  // The slices are copied; only referenced memory is borrowed, owned by `owner` until written.
  BorrowedData(const DataSlice* dataSlices, size_t dataSliceCount, std::shared_ptr<void> owner)
      : owner_{std::move(owner)} {
    segments_.reserve(dataSliceCount);
    for (size_t i = 0; i < dataSliceCount; ++i) {
      segments_.push_back(dataSlices[i]);
      size_ += dataSlices[i].size;
    }
  }
 
  const std::vector<DataSlice>& getSegments() const {
    return segments_;
  }
 
  size_t getDataSize() const {
    return size_;
  }
 
 private:
  std::vector<DataSlice> segments_;
  std::shared_ptr<void> owner_;
  size_t size_{0};
};
 
} // namespace vrs