protobuf/rust/cpp_kernel/map.cc

#include "rust/cpp_kernel/map.h"

#include <cstddef>
#include <cstdint>
#include <string>
#include <type_traits>
#include <utility>

#include "absl/log/absl_log.h"
#include "google/protobuf/map.h"
#include "google/protobuf/message.h"
#include "google/protobuf/message_lite.h"
#include "rust/cpp_kernel/strings.h"

namespace google {
namespace protobuf {
namespace rust {
namespace {

template <typename T>
struct FromViewType {
  using type = T;
};

template <>
struct FromViewType<PtrAndLen> {
  using type = std::string;
};

template <typename Key>
using KeyMap = internal::KeyMapBase<
    internal::KeyForBase<typename FromViewType<Key>::type>>;

internal::MapNodeSizeInfoT GetSizeInfo(size_t key_size,
                                       const google::protobuf::MessageLite* value) {
  // Each map node consists of a NodeBase followed by a std::pair<Key, Value>.
  // We need to compute the offset of the value and the total size of the node.
  size_t node_and_key_size = sizeof(internal::NodeBase) + key_size;
  uint16_t value_size;
  uint8_t value_alignment;
  internal::RustMapHelper::GetSizeAndAlignment(value, &value_size,
                                               &value_alignment);
  // Round node_and_key_size up to the nearest multiple of value_alignment.
  uint16_t offset =
      (((node_and_key_size - 1) / value_alignment) + 1) * value_alignment;
  return internal::RustMapHelper::MakeSizeInfo(offset + value_size, offset);
}

template <typename Key>
internal::MapNodeSizeInfoT GetSizeInfo(const google::protobuf::MessageLite* value) {
  return GetSizeInfo(sizeof(Key), value);
}

template <typename Key>
void DestroyMapNode(internal::UntypedMapBase* m, internal::NodeBase* node,
                    internal::MapNodeSizeInfoT size_info) {
  if constexpr (std::is_same<Key, PtrAndLen>::value) {
    static_cast<std::string*>(node->GetVoidKey())->~basic_string();
  }
  internal::RustMapHelper::DestroyMessage(
      static_cast<MessageLite*>(node->GetVoidValue(size_info)));
  internal::RustMapHelper::DeallocNode(m, node, size_info);
}

template <typename Key>
bool Insert(internal::UntypedMapBase* m, Key key, MessageLite* value) {
  internal::MapNodeSizeInfoT size_info =
      GetSizeInfo<typename FromViewType<Key>::type>(value);
  internal::NodeBase* node = internal::RustMapHelper::AllocNode(m, size_info);
  if constexpr (std::is_same<Key, PtrAndLen>::value) {
    new (node->GetVoidKey()) std::string(key.ptr, key.len);
  } else {
    *static_cast<Key*>(node->GetVoidKey()) = key;
  }

  MessageLite* new_msg = internal::RustMapHelper::PlacementNew(
      value, node->GetVoidValue(size_info));
  auto* full_msg = DynamicCastMessage<Message>(new_msg);

  // If we are working with a full (non-lite) proto, we reflectively swap the
  // value into place. Otherwise, we have to perform a copy.
  if (full_msg != nullptr) {
    full_msg->GetReflection()->Swap(full_msg,
                                    DynamicCastMessage<Message>(value));
  } else {
    new_msg->CheckTypeAndMergeFrom(*value);
  }

  node = internal::RustMapHelper::InsertOrReplaceNode(
      static_cast<KeyMap<Key>*>(m), node);
  if (node == nullptr) {
    return true;
  }
  DestroyMapNode<Key>(m, node, size_info);
  return false;
}

template <typename Map, typename Key,
          typename = typename std::enable_if<
              !std::is_same<Key, google::protobuf::rust::PtrAndLen>::value>::type>
internal::RustMapHelper::NodeAndBucket FindHelper(Map* m, Key key) {
  return internal::RustMapHelper::FindHelper(
      m, static_cast<internal::KeyForBase<Key>>(key));
}

template <typename Map>
internal::RustMapHelper::NodeAndBucket FindHelper(Map* m,
                                                  google::protobuf::rust::PtrAndLen key) {
  return internal::RustMapHelper::FindHelper(
      m, absl::string_view(key.ptr, key.len));
}

template <typename Key>
bool Get(internal::UntypedMapBase* m, const google::protobuf::MessageLite* prototype,
         Key key, MessageLite** value) {
  internal::MapNodeSizeInfoT size_info =
      GetSizeInfo<typename FromViewType<Key>::type>(prototype);
  auto* map_base = static_cast<KeyMap<Key>*>(m);
  internal::RustMapHelper::NodeAndBucket result = FindHelper(map_base, key);
  if (result.node == nullptr) {
    return false;
  }
  *value = static_cast<MessageLite*>(result.node->GetVoidValue(size_info));
  return true;
}

template <typename Key>
bool Remove(internal::UntypedMapBase* m, const google::protobuf::MessageLite* prototype,
            Key key) {
  internal::MapNodeSizeInfoT size_info =
      GetSizeInfo<typename FromViewType<Key>::type>(prototype);
  auto* map_base = static_cast<KeyMap<Key>*>(m);
  internal::RustMapHelper::NodeAndBucket result = FindHelper(map_base, key);
  if (result.node == nullptr) {
    return false;
  }
  internal::RustMapHelper::EraseNoDestroy(map_base, result.bucket, result.node);
  DestroyMapNode<Key>(m, result.node, size_info);
  return true;
}

template <typename Key>
void IterGet(const internal::UntypedMapIterator* iter,
             const google::protobuf::MessageLite* prototype, Key* key,
             MessageLite** value) {
  internal::MapNodeSizeInfoT size_info =
      GetSizeInfo<typename FromViewType<Key>::type>(prototype);
  internal::NodeBase* node = iter->node_;
  if constexpr (std::is_same<Key, PtrAndLen>::value) {
    const std::string* s = static_cast<const std::string*>(node->GetVoidKey());
    *key = PtrAndLen{s->data(), s->size()};
  } else {
    *key = *static_cast<const Key*>(node->GetVoidKey());
  }
  *value = static_cast<MessageLite*>(node->GetVoidValue(size_info));
}

// LINT.IfChange(map_key_category)
enum class MapKeyCategory : uint8_t {
  kOneByte = 0,
  kFourBytes = 1,
  kEightBytes = 2,
  kStdString = 3,
};
// LINT.ThenChange(//depot/google3/third_party/protobuf/rust/cpp.rs:map_key_category)

size_t KeySize(MapKeyCategory category) {
  switch (category) {
    case MapKeyCategory::kOneByte:
      return 1;
    case MapKeyCategory::kFourBytes:
      return 4;
    case MapKeyCategory::kEightBytes:
      return 8;
    case MapKeyCategory::kStdString:
      return sizeof(std::string);
    default:
      ABSL_DLOG(FATAL) << "Unexpected value of MapKeyCategory enum";
  }
}

void ClearMap(internal::UntypedMapBase* m, MapKeyCategory category,
              bool reset_table, const google::protobuf::MessageLite* prototype) {
  internal::MapNodeSizeInfoT size_info =
      GetSizeInfo(KeySize(category), prototype);
  if (internal::RustMapHelper::IsGlobalEmptyTable(m)) return;
  uint8_t bits = internal::RustMapHelper::kValueIsProto;
  if (category == MapKeyCategory::kStdString) {
    bits |= internal::RustMapHelper::kKeyIsString;
  }
  internal::RustMapHelper::ClearTable(
      m, internal::RustMapHelper::ClearInput{size_info, bits, reset_table,
                                             /* destroy_node = */ nullptr});
}

}  // namespace
}  // namespace rust
}  // namespace protobuf
}  // namespace google

extern "C" {

void proto2_rust_thunk_UntypedMapIterator_increment(
    google::protobuf::internal::UntypedMapIterator* iter) {
  iter->PlusPlus();
}

google::protobuf::internal::UntypedMapBase* proto2_rust_map_new() {
  return new google::protobuf::internal::UntypedMapBase(/* arena = */ nullptr);
}

void proto2_rust_map_free(google::protobuf::internal::UntypedMapBase* m,
                          google::protobuf::rust::MapKeyCategory category,
                          const google::protobuf::MessageLite* prototype) {
  google::protobuf::rust::ClearMap(m, category, /* reset_table = */ false, prototype);
  delete m;
}

void proto2_rust_map_clear(google::protobuf::internal::UntypedMapBase* m,
                           google::protobuf::rust::MapKeyCategory category,
                           const google::protobuf::MessageLite* prototype) {
  google::protobuf::rust::ClearMap(m, category, /* reset_table = */ true, prototype);
}

size_t proto2_rust_map_size(google::protobuf::internal::UntypedMapBase* m) {
  return m->size();
}

google::protobuf::internal::UntypedMapIterator proto2_rust_map_iter(
    google::protobuf::internal::UntypedMapBase* m) {
  return m->begin();
}

#define DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(cpp_type, suffix)                 \
  bool proto2_rust_map_insert_##suffix(google::protobuf::internal::UntypedMapBase* m,  \
                                       cpp_type key,                         \
                                       google::protobuf::MessageLite* value) {         \
    return google::protobuf::rust::Insert(m, key, value);                              \
  }                                                                          \
                                                                             \
  bool proto2_rust_map_get_##suffix(google::protobuf::internal::UntypedMapBase* m,     \
                                    const google::protobuf::MessageLite* prototype,    \
                                    cpp_type key,                            \
                                    google::protobuf::MessageLite** value) {           \
    return google::protobuf::rust::Get(m, prototype, key, value);                      \
  }                                                                          \
                                                                             \
  bool proto2_rust_map_remove_##suffix(google::protobuf::internal::UntypedMapBase* m,  \
                                       const google::protobuf::MessageLite* prototype, \
                                       cpp_type key) {                       \
    return google::protobuf::rust::Remove(m, prototype, key);                          \
  }                                                                          \
                                                                             \
  void proto2_rust_map_iter_get_##suffix(                                    \
      const google::protobuf::internal::UntypedMapIterator* iter,                      \
      const google::protobuf::MessageLite* prototype, cpp_type* key,                   \
      google::protobuf::MessageLite** value) {                                         \
    return google::protobuf::rust::IterGet(iter, prototype, key, value);               \
  }

DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(int32_t, i32)
DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(uint32_t, u32)
DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(int64_t, i64)
DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(uint64_t, u64)
DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(bool, bool)
DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(google::protobuf::rust::PtrAndLen, ProtoString)

__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(int32_t, i32, int32_t,
                                                   int32_t, value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(uint32_t, u32, uint32_t,
                                                   uint32_t, value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(float, f32, float, float,
                                                   value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(double, f64, double, double,
                                                   value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(bool, bool, bool, bool,
                                                   value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(uint64_t, u64, uint64_t,
                                                   uint64_t, value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(int64_t, i64, int64_t,
                                                   int64_t, value, cpp_value);
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(
    std::string, ProtoBytes, google::protobuf::rust::PtrAndLen, std::string*,
    std::move(*value),
    (google::protobuf::rust::PtrAndLen{cpp_value.data(), cpp_value.size()}));
__PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(
    std::string, ProtoString, google::protobuf::rust::PtrAndLen, std::string*,
    std::move(*value),
    (google::protobuf::rust::PtrAndLen{cpp_value.data(), cpp_value.size()}));

}  // extern "C"