DataPieceVector.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
 
#ifndef DATA_PIECES_VECTOR_H
#define DATA_PIECES_VECTOR_H
 
#include <type_traits>
 
#ifndef DATA_PIECES_H
#include <vrs/DataPieces.h>
#endif
 
namespace vrs {
 
using std::ostream;
using std::string;
using std::vector;
 
template <typename T>
class DataPieceVector : public DataPiece {
 public:
  explicit DataPieceVector(string label, vector<T> defaultValues = {})
      : DataPiece(std::move(label), DataPieceType::Vector, DataLayout::kVariableSize),
        defaultValues_{std::move(defaultValues)} {}
  explicit DataPieceVector(const MakerBundle& bundle);
 
  const string& getElementTypeName() const override {
    return vrs::getTypeName<T>();
  }
  size_t getVariableSize() const override {
    return sizeof(T) * stagedValues_.size();
  }
  size_t collectVariableData(int8_t* data, size_t bufferSize) override {
    const size_t writtenSize = min(bufferSize, getVariableSize());
    if (writtenSize > 0) {
      unalignedCopy(data, stagedValues_.data(), writtenSize);
    }
    return writtenSize;
  }
 
  const vector<T>& stagedValues() const {
    return stagedValues_;
  }
  vector<T>& stagedValues() {
    return stagedValues_;
  }
  void stage(const vector<T>& values) {
    stagedValues_ = values;
  }
  void stage(vector<T>&& values) {
    stagedValues_ = std::move(values);
  }
 
  void stage(const T* values, size_t count) {
    stagedValues_.resize(count);
    if (count > 0) {
      unalignedCopy(stagedValues_.data(), values, count * sizeof(T));
    }
  }
  template <size_t n>
  void stage(const T (&arr)[n]) {
    stage(arr, n);
  }
 
  bool get(vector<T>& outValues) const {
    size_t count = 0;
    const T* ptr = layout_.getVarData<T>(offset_, count);
    if (ptr != nullptr && count > 0) {
      outValues.resize(count);
      unalignedCopy(outValues.data(), ptr, count * sizeof(T));
      return true;
    }
    if (pieceIndex_ != DataLayout::kNotFound) {
      outValues.clear();
      return true;
    }
    outValues = defaultValues_;
    return false;
  }
  template <
      typename U = T,
      std::enable_if_t<std::is_same<U, T>::value && std::is_trivially_copyable<U>::value, int> = 0>
  bool getRawData(const void*& outData, size_t& outCount) const {
    outData = layout_.getVarData<T>(offset_, outCount);
    if (outData != nullptr && outCount > 0) {
      return true;
    }
    outData = nullptr;
    outCount = 0;
    return pieceIndex_ != DataLayout::kNotFound;
  }
 
  void setDefault(const T* defaultValues, size_t count) {
    defaultValues_.resize(count);
    if (count > 0) {
      unalignedCopy(defaultValues_.data(), defaultValues, count * sizeof(T));
    }
  }
  template <size_t n>
  void setDefault(const T (&arr)[n]) {
    setDefault(arr, n);
  }
  void setDefault(const vector<T>& values) {
    defaultValues_ = values;
  }
  void setDefault(vector<T>&& values) {
    defaultValues_ = std::move(values);
  }
  const vector<T>& getDefault() const {
    return defaultValues_;
  }
  void initToDefault() override {
    stagedValues_ = defaultValues_;
  }
 
  bool patchValue(const vector<T>& values) const {
    auto* patchedPiece = layout_.getMappedPiece<DataPieceVector<T>>(pieceIndex_);
    return patchedPiece != nullptr && (patchedPiece->stage(values), true);
  }
 
  bool isAvailable() const override {
    size_t count = 0;
    return layout_.getVarData<T>(offset_, count) != nullptr;
  }
 
  void print(ostream& out, const string& indent) const override;
  void printCompact(ostream& out, const string& indent) const override;
 
  bool isSame(const DataPiece* rhs) const override;
 
  void serialize(JsonWrapper& jsonWrapper, const JsonFormatProfileSpec& profile) override;
 
  bool stageCurrentValue() override {
    return get(stagedValues());
  }
 
  unique_ptr<DataPiece> clone() const override {
    auto other = std::make_unique<DataPieceVector<T>>(getLabel());
    other->tags_ = tags_;
    other->required_ = required_;
    other->defaultValues_ = defaultValues_;
    return other;
  }
 
 protected:
  bool copyFrom(const DataPiece* original) override {
    const DataPieceVector<T>* originalVect = reinterpret_cast<const DataPieceVector<T>*>(original);
    return originalVect->get(stagedValues_);
  }
 
 private:
  vector<T> stagedValues_; // values to write to disk
  vector<T> defaultValues_;
};
 
// Specializations declaration, required here for come compilers.
template <>
VRS_API void DataPieceVector<string>::stage(const string* values, size_t count);
template <>
VRS_API size_t DataPieceVector<string>::getVariableSize() const;
template <>
VRS_API size_t DataPieceVector<string>::collectVariableData(int8_t* data, size_t bufferSize);
template <>
void DataPieceVector<string>::setDefault(const string* defaultValues, size_t count) = delete;
template <>
VRS_API bool DataPieceVector<string>::get(vector<string>& outValues) const;
template <>
VRS_API void DataPieceVector<string>::print(ostream& out, const string& indent) const;
template <>
VRS_API void DataPieceVector<string>::printCompact(ostream& out, const string& indent) const;
 
} // namespace vrs
 
#endif // DATA_PIECES_VECTOR_H