// Protocol Buffers - Google's data interchange format
// Copyright 2023 Google LLC.  All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

/* Test of mini table accessors.
 *
 * Messages are created and mutated using generated code, and then
 * accessed through reflective APIs exposed through mini table accessors.
 */

#include "upb/message/copy.h"

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <string>
#include <string_view>
#include <vector>

#include <gtest/gtest.h>
#include "google/protobuf/test_messages_proto2.upb.h"
#include "google/protobuf/test_messages_proto2.upb_minitable.h"
#include "upb/base/string_view.h"
#include "upb/base/upcast.h"
#include "upb/mem/arena.h"
#include "upb/message/accessors.h"
#include "upb/message/internal/accessors.h"
#include "upb/message/internal/extension.h"
#include "upb/message/internal/message.h"
#include "upb/message/map.h"
#include "upb/message/message.h"
#include "upb/message/unknown_fields.h"
#include "upb/mini_table/field.h"
#include "upb/mini_table/message.h"
#include "upb/test/test.upb.h"
#include "upb/test/test.upb_minitable.h"
#include "upb/wire/encode.h"

// Must be last.
#include "upb/port/def.inc"

namespace {

// Proto2 test messages field numbers used for reflective access.
const uint32_t kFieldOptionalInt32 = 1;
const uint32_t kFieldOptionalString = 14;
const uint32_t kFieldOptionalNestedMessage = 18;

const char kTestStr1[] = "Hello1";
const char kTestStr2[] = "HelloWorld2";
const int32_t kTestInt32 = 567;
const int32_t kTestNestedInt32 = 123;
const int32_t kTestNestedInt64 = 123456789;

const upb_MiniTableField* find_proto2_field(int field_number) {
  return upb_MiniTable_FindFieldByNumber(
      &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
      field_number);
}

TEST(GeneratedCode, DeepCloneMessageScalarAndString) {
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);
  const upb_MiniTableField* optional_int32_field =
      find_proto2_field(kFieldOptionalInt32);
  const upb_MiniTableField* optional_string_field =
      find_proto2_field(kFieldOptionalString);
  ASSERT_TRUE(upb_Message_SetInt32(UPB_UPCAST(msg), optional_int32_field,
                                   kTestInt32, nullptr));
  char* string_in_arena =
      (char*)upb_Arena_Malloc(source_arena, sizeof(kTestStr1));
  memcpy(string_in_arena, kTestStr1, sizeof(kTestStr1));
  ASSERT_TRUE(upb_Message_SetString(
      UPB_UPCAST(msg), optional_string_field,
      upb_StringView_FromDataAndSize(string_in_arena, sizeof(kTestStr1) - 1),
      source_arena));
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)upb_Message_DeepClone(
          UPB_UPCAST(msg),
          &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
          arena);
  // After cloning overwrite values and destroy source arena for MSAN.
  memset(string_in_arena, 0, sizeof(kTestStr1));
  upb_Arena_Free(source_arena);
  EXPECT_TRUE(
      upb_Message_HasBaseField(UPB_UPCAST(clone), optional_int32_field));
  EXPECT_EQ(upb_Message_GetInt32(UPB_UPCAST(clone), optional_int32_field, 0),
            kTestInt32);
  EXPECT_TRUE(
      upb_Message_HasBaseField(UPB_UPCAST(clone), optional_string_field));
  EXPECT_EQ(upb_Message_GetString(UPB_UPCAST(clone), optional_string_field,
                                  upb_StringView_FromDataAndSize(nullptr, 0))
                .size,
            sizeof(kTestStr1) - 1);
  EXPECT_TRUE(upb_StringView_IsEqual(
      upb_Message_GetString(UPB_UPCAST(clone), optional_string_field,
                            upb_StringView_FromDataAndSize(nullptr, 0)),
      upb_StringView_FromString(kTestStr1)));
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageSubMessage) {
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);
  const upb_MiniTableField* nested_message_field =
      find_proto2_field(kFieldOptionalNestedMessage);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage* nested =
      protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_new(
          source_arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_set_a(
      nested, kTestNestedInt32);
  upb_Message_SetMessage(UPB_UPCAST(msg), nested_message_field,
                         UPB_UPCAST(nested));
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)upb_Message_DeepClone(
          UPB_UPCAST(msg),
          &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
          arena);
  // After cloning overwrite values and destroy source arena for MSAN.
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_set_a(nested,
                                                                       0);
  upb_Arena_Free(source_arena);
  EXPECT_TRUE(
      upb_Message_HasBaseField(UPB_UPCAST(clone), nested_message_field));
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage*
      cloned_nested =
          (protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage*)
              upb_Message_GetMessage(UPB_UPCAST(clone), nested_message_field);
  EXPECT_EQ(protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_a(
                cloned_nested),
            kTestNestedInt32);
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageArrayField) {
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);
  std::vector<int32_t> array_test_values = {3, 4, 5};
  for (int32_t value : array_test_values) {
    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_add_repeated_int32(
            msg, value, source_arena));
  }
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)upb_Message_DeepClone(
          UPB_UPCAST(msg),
          &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
          arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_clear_repeated_sint32(msg);
  upb_Arena_Free(source_arena);
  size_t cloned_size = 0;
  const int32_t* cloned_values =
      protobuf_test_messages_proto2_TestAllTypesProto2_repeated_int32(
          clone, &cloned_size);
  EXPECT_EQ(cloned_size, array_test_values.size());
  int index = 0;
  for (int32_t value : array_test_values) {
    EXPECT_EQ(cloned_values[index++], value);
  }
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageMapField) {
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_bool_set(
          msg, 0, true, source_arena));
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_double_set(
          msg, 12, 1200.5, source_arena));
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_string_string_set(
          msg, upb_StringView_FromString("key1"),
          upb_StringView_FromString("value1"), source_arena));
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage* nested =
      protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_new(
          source_arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_set_a(
      nested, kTestNestedInt32);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage* nested2 =
      protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_new(
          source_arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_set_a(
      nested2, kTestNestedInt64);
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_nested_message_set(
          msg, 1, nested2, source_arena));
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_string_nested_message_set(
          msg, upb_StringView_FromString("nestedkey1"), nested, source_arena));

  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)upb_Message_DeepClone(
          UPB_UPCAST(msg),
          &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
          arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_set_a(nested,
                                                                       0);
  upb_Arena_Free(source_arena);
  // Test map<int32, bool>.
  {
    int32_t key;
    bool value;
    size_t iter = kUpb_Map_Begin;

    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_bool_next(
            clone, &key, &value, &iter));
    EXPECT_EQ(key, 0);
    EXPECT_EQ(value, true);

    ASSERT_FALSE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_bool_next(
            clone, &key, &value, &iter));
  }

  // Test map<int32, double>.
  {
    int32_t key;
    double value;
    size_t iter = kUpb_Map_Begin;

    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_double_next(
            clone, &key, &value, &iter));
    EXPECT_EQ(key, 12);
    EXPECT_EQ(value, 1200.5);

    ASSERT_FALSE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_double_next(
            clone, &key, &value, &iter));
  }

  // Test map<int32, NestedMessage>.
  {
    int32_t key;
    const protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage* value;
    size_t iter = kUpb_Map_Begin;
    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_nested_message_next(
            clone, &key, &value, &iter));
    EXPECT_EQ(key, 1);
    ASSERT_NE(value, nullptr);
    EXPECT_EQ(
        protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_a(value),
        kTestNestedInt64);
    ASSERT_FALSE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_nested_message_next(
            clone, &key, &value, &iter));
  }

  // Test map<string, string>.
  {
    upb_StringView key;
    upb_StringView value;
    size_t iter = kUpb_Map_Begin;

    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_string_string_next(
            clone, &key, &value, &iter));
    EXPECT_TRUE(upb_StringView_IsEqual(key, upb_StringView_FromString("key1")));
    EXPECT_TRUE(
        upb_StringView_IsEqual(value, upb_StringView_FromString("value1")));
  }

  // Test map<string, NestedMessage>.
  {
    upb_StringView key;
    const protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage* value;
    size_t iter = kUpb_Map_Begin;
    ASSERT_TRUE(
        protobuf_test_messages_proto2_TestAllTypesProto2_map_string_nested_message_next(
            clone, &key, &value, &iter));
    EXPECT_TRUE(
        upb_StringView_IsEqual(key, upb_StringView_FromString("nestedkey1")));
    ASSERT_NE(value, nullptr);
    EXPECT_EQ(
        protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_a(value),
        kTestNestedInt32);
  }

  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageExtensions) {
  // Alloc and fill in test message with extension.
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrect* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrect_new(
          source_arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1*
      ext1 =
          protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_new(
              source_arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_set_str(
      ext1, upb_StringView_FromString(kTestStr1));
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_set_message_set_extension(
      msg, ext1, source_arena);
  // Create clone.
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrect* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrect*)
          upb_Message_DeepClone(
              UPB_UPCAST(msg),
              &protobuf_0test_0messages__proto2__TestAllTypesProto2__MessageSetCorrect_msg_init,
              arena);

  // Mutate original extension.
  protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_set_str(
      ext1, upb_StringView_FromString(kTestStr2));
  upb_Arena_Free(source_arena);

  const protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1*
      cloned_ext =
          protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_message_set_extension(
              clone);
  ASSERT_NE(cloned_ext, nullptr);
  EXPECT_TRUE(upb_StringView_IsEqual(
      protobuf_test_messages_proto2_TestAllTypesProto2_MessageSetCorrectExtension1_str(
          cloned_ext),
      upb_StringView_FromString(kTestStr1)));
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageWithUnknowns) {
  upb_Arena* source_arena = upb_Arena_New();
  upb_Arena* unknown_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_int32_double_set(
          msg, 12, 1200.5, source_arena));
  ASSERT_TRUE(
      protobuf_test_messages_proto2_TestAllTypesProto2_map_string_string_set(
          msg, upb_StringView_FromString("key1"),
          upb_StringView_FromString("value1"), source_arena));
  // Create unknown data.
  protobuf_test_messages_proto2_UnknownToTestAllTypes* unknown_source =
      protobuf_test_messages_proto2_UnknownToTestAllTypes_new(unknown_arena);
  protobuf_test_messages_proto2_UnknownToTestAllTypes_set_optional_bool(
      unknown_source, true);
  protobuf_test_messages_proto2_UnknownToTestAllTypes_set_optional_int32(
      unknown_source, 123);
  // Encode unknown message to bytes.
  size_t len;
  char* data;
  upb_Arena* encode_arena = upb_Arena_New();
  upb_EncodeStatus status = upb_Encode(
      UPB_UPCAST(unknown_source),
      &protobuf_0test_0messages__proto2__UnknownToTestAllTypes_msg_init,
      kUpb_EncodeOption_CheckRequired, encode_arena, &data, &len);
  ASSERT_EQ(status, kUpb_EncodeStatus_Ok);
  std::string unknown_data(data, len);
  // Add unknown data.
  ASSERT_TRUE(UPB_PRIVATE(_upb_Message_AddUnknown)(
      UPB_UPCAST(msg), data, len, source_arena, kUpb_AddUnknown_Copy));
  // Create clone.
  upb_Arena* clone_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)upb_Message_DeepClone(
          UPB_UPCAST(msg),
          &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
          clone_arena);
  upb_Arena_Free(source_arena);
  upb_Arena_Free(unknown_arena);
  upb_Arena_Free(encode_arena);
  // Read unknown data from clone and verify.
  std::string cloned_unknown_data;
  upb_MessageUnknown unknown;
  uintptr_t iter = kUpb_Message_UnknownBegin;
  while (upb_Message_NextUnknown2(UPB_UPCAST(clone), &unknown, &iter)) {
    ASSERT_EQ(unknown.type, kUpb_MessageUnknownType_StringView);
    cloned_unknown_data.append(unknown.value.bytes.data,
                               unknown.value.bytes.size);
  }
  EXPECT_EQ(unknown_data, cloned_unknown_data);
  upb_Arena_Free(clone_arena);
}

TEST(GeneratedCode, ShallowCopyMessage) {
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_set_optional_int32(
      msg, kTestInt32);
  char* string_in_arena = (char*)upb_Arena_Malloc(arena, sizeof(kTestStr1));
  memcpy(string_in_arena, kTestStr1, sizeof(kTestStr1));
  protobuf_test_messages_proto2_TestAllTypesProto2_set_optional_string(
      msg,
      upb_StringView_FromDataAndSize(string_in_arena, sizeof(kTestStr1) - 1));

  protobuf_test_messages_proto2_TestAllTypesProto2* dst =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);
  EXPECT_TRUE(upb_Message_ShallowCopy(
      UPB_UPCAST(dst), UPB_UPCAST(msg),
      &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init, arena));

  EXPECT_EQ(
      protobuf_test_messages_proto2_TestAllTypesProto2_optional_int32(dst),
      kTestInt32);
  upb_StringView dst_str =
      protobuf_test_messages_proto2_TestAllTypesProto2_optional_string(dst);
  EXPECT_EQ(dst_str.size, sizeof(kTestStr1) - 1);
  EXPECT_EQ(dst_str.data, string_in_arena);

  upb_Arena_Free(arena);
}

TEST(GeneratedCode, ShallowCopyIncludesExtensions) {
  upb_Arena* source_arena = upb_Arena_New();
  upb_test_ModelWithExtensions* msg =
      upb_test_ModelWithExtensions_new(source_arena);
  upb_test_ModelExtension1* ext1 = upb_test_ModelExtension1_new(source_arena);
  upb_test_ModelExtension1_set_str(ext1, upb_StringView_FromString(kTestStr1));
  upb_test_ModelExtension1_set_model_ext(msg, ext1, source_arena);

  upb_Arena* arena = upb_Arena_New();
  upb_test_ModelWithExtensions* dst = upb_test_ModelWithExtensions_new(arena);

  EXPECT_TRUE(upb_Message_ShallowCopy(UPB_UPCAST(dst), UPB_UPCAST(msg),
                                      &upb_0test__ModelWithExtensions_msg_init,
                                      arena));

  // Modifying extension in dst should NOT affect msg because extensions are
  // copied.
  upb_test_ModelExtension1* ext2 = upb_test_ModelExtension1_new(arena);
  upb_test_ModelExtension1_set_str(ext2, upb_StringView_FromString(kTestStr2));

  upb_test_ModelExtension1_set_model_ext(dst, ext2, arena);

  const upb_test_ModelExtension1* src_ext =
      upb_test_ModelExtension1_model_ext(msg);

  EXPECT_EQ(UPB_UPCAST(src_ext), UPB_UPCAST(ext1));

  const upb_test_ModelExtension1* dst_ext =
      upb_test_ModelExtension1_model_ext(dst);

  EXPECT_EQ(UPB_UPCAST(dst_ext), UPB_UPCAST(ext2));

  upb_Arena_Free(source_arena);
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, ShallowCopyHandlesClearedExtensions) {
  upb_Arena* source_arena = upb_Arena_New();
  upb_test_ModelWithExtensions* msg =
      upb_test_ModelWithExtensions_new(source_arena);
  upb_test_ModelExtension1* ext1 = upb_test_ModelExtension1_new(source_arena);
  upb_test_ModelExtension1_set_str(ext1, upb_StringView_FromString(kTestStr1));
  upb_test_ModelExtension1_set_model_ext(msg, ext1, source_arena);

  // Clear the extension, creating a tombstone in aux_data.
  upb_Message_ClearExtension(UPB_UPCAST(msg),
                             upb_test_ModelExtension1_model_ext_ext);

  // Verify source has a tombstone (size is 1).
  upb_Message_Internal* msg_in =
      UPB_PRIVATE(_upb_Message_GetInternal)(UPB_UPCAST(msg));
  EXPECT_NE(msg_in, nullptr);
  EXPECT_EQ(msg_in->size, 1);

  upb_Arena* arena = upb_Arena_New();
  upb_test_ModelWithExtensions* dst = upb_test_ModelWithExtensions_new(arena);

  // This should not crash.
  EXPECT_TRUE(upb_Message_ShallowCopy(UPB_UPCAST(dst), UPB_UPCAST(msg),
                                      &upb_0test__ModelWithExtensions_msg_init,
                                      arena));

  // The extension should not be present in dst.
  EXPECT_EQ(upb_test_ModelExtension1_model_ext(dst), nullptr);

  // Verify internal array size is compacted in dst (size is 0).
  upb_Message_Internal* dst_in =
      UPB_PRIVATE(_upb_Message_GetInternal)(UPB_UPCAST(dst));
  EXPECT_NE(dst_in, nullptr);
  EXPECT_EQ(dst_in->size, 0);

  upb_Arena_Free(source_arena);
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneHandlesClearedExtensions) {
  upb_Arena* source_arena = upb_Arena_New();
  upb_test_ModelWithExtensions* msg =
      upb_test_ModelWithExtensions_new(source_arena);
  upb_test_ModelExtension1* ext1 = upb_test_ModelExtension1_new(source_arena);
  upb_test_ModelExtension1_set_str(ext1, upb_StringView_FromString(kTestStr1));
  upb_test_ModelExtension1_set_model_ext(msg, ext1, source_arena);

  // Clear the extension, creating a tombstone in aux_data.
  upb_Message_ClearExtension(UPB_UPCAST(msg),
                             upb_test_ModelExtension1_model_ext_ext);

  // Verify source has a tombstone (size is 1).
  upb_Message_Internal* msg_in =
      UPB_PRIVATE(_upb_Message_GetInternal)(UPB_UPCAST(msg));
  EXPECT_NE(msg_in, nullptr);
  EXPECT_EQ(msg_in->size, 1);

  upb_Arena* arena = upb_Arena_New();

  // This should not crash.
  upb_Message* clone = upb_Message_DeepClone(
      UPB_UPCAST(msg), &upb_0test__ModelWithExtensions_msg_init, arena);
  EXPECT_NE(clone, nullptr);

  // The extension should not be present in clone.
  EXPECT_EQ(
      upb_test_ModelExtension1_model_ext((upb_test_ModelWithExtensions*)clone),
      nullptr);

  // Verify internal array is compacted (DeepClone doesn't create internal data
  // if empty).
  upb_Message_Internal* clone_in = UPB_PRIVATE(_upb_Message_GetInternal)(clone);
  EXPECT_EQ(clone_in, nullptr);

  upb_Arena_Free(source_arena);
  upb_Arena_Free(arena);
}

std::vector<std::string_view> GetUnknownFields(const upb_Message* msg) {
  std::vector<std::string_view> result;
  upb_MessageUnknown data;
  uintptr_t iter = kUpb_Message_UnknownBegin;
  while (upb_Message_NextUnknown2(msg, &data, &iter)) {
    EXPECT_EQ(data.type, kUpb_MessageUnknownType_StringView);
    result.push_back(
        std::string_view(data.value.bytes.data, data.value.bytes.size));
  }
  return result;
}

TEST(GeneratedCode, ShallowCopyIncludesUnknowns) {
  upb_Arena* source_arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(source_arena);

  // Add some unknown data.
  std::string_view unknown_data1 =
      "\x08\x01";  // Field 1, wire type 0 (varint), value 1
  ASSERT_TRUE(UPB_PRIVATE(_upb_Message_AddUnknown)(
      UPB_UPCAST(msg), unknown_data1.data(), unknown_data1.size(), source_arena,
      kUpb_AddUnknown_Copy));
  std::vector<std::string_view> only_unknown_data1 = {unknown_data1};

  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* dst =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);

  EXPECT_TRUE(upb_Message_ShallowCopy(
      UPB_UPCAST(dst), UPB_UPCAST(msg),
      &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init, arena));

  // Verify dst has the unknown data.
  EXPECT_EQ(GetUnknownFields(UPB_UPCAST(dst)), only_unknown_data1);

  // Modify the unknown data view in dst to ensure it's a separate view.
  // Use upb_Message_DeleteUnknown to delete a trailing part.
  upb_MessageUnknown dst_unknown;
  uintptr_t iter = kUpb_Message_UnknownBegin;
  EXPECT_TRUE(upb_Message_NextUnknown2(UPB_UPCAST(dst), &dst_unknown, &iter));
  ASSERT_EQ(dst_unknown.type, kUpb_MessageUnknownType_StringView);
  upb_StringView to_delete;
  to_delete.size = 1;
  to_delete.data = dst_unknown.value.bytes.data + dst_unknown.value.bytes.size -
                   to_delete.size;
  upb_Message_DeleteUnknownStatus status =
      upb_Message_DeleteUnknown(UPB_UPCAST(dst), &to_delete, &iter, arena);
  EXPECT_EQ(status, kUpb_DeleteUnknown_DeletedLast);

  // Verify src still has the original data and size.
  EXPECT_EQ(GetUnknownFields(UPB_UPCAST(msg)), only_unknown_data1);

  // Add MORE unknown data to dst.
  std::string_view unknown_data2 =
      "\x10\x02";  // Field 2, wire type 0 (varint), value 2
  ASSERT_TRUE(UPB_PRIVATE(_upb_Message_AddUnknown)(
      UPB_UPCAST(dst), unknown_data2.data(), unknown_data2.size(), arena,
      kUpb_AddUnknown_Copy));

  // Verify src still only has the first unknown data.
  EXPECT_EQ(GetUnknownFields(UPB_UPCAST(msg)), only_unknown_data1);

  upb_Arena_Free(source_arena);
  upb_Arena_Free(arena);
}

TEST(GeneratedCode, ShallowCloneMessage) {
  upb_Arena* arena = upb_Arena_New();
  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);
  protobuf_test_messages_proto2_TestAllTypesProto2_set_optional_int32(
      msg, kTestInt32);

  protobuf_test_messages_proto2_TestAllTypesProto2* clone =
      (protobuf_test_messages_proto2_TestAllTypesProto2*)
          upb_Message_ShallowClone(
              UPB_UPCAST(msg),
              &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init,
              arena);

  EXPECT_EQ(
      protobuf_test_messages_proto2_TestAllTypesProto2_optional_int32(clone),
      kTestInt32);

  upb_Arena_Free(arena);
}

TEST(GeneratedCode, DeepCloneMessageNonCanonicalExtensions) {
  upb_Arena* source_arena = upb_Arena_New();
  upb_test_ModelWithExtensions* msg =
      upb_test_ModelWithExtensions_new(source_arena);
  upb_test_ModelExtension1* ext1 = upb_test_ModelExtension1_new(source_arena);
  upb_test_ModelExtension1_set_str(ext1,
                                   upb_StringView_FromString("LifecycleValue"));

  // Attach as non-canonical extension
  EXPECT_TRUE(UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)(
      UPB_UPCAST(msg), upb_test_ModelExtension1_model_ext_ext, &ext1,
      source_arena));

  // Deep clone msg to clone
  upb_Arena* arena = upb_Arena_New();
  upb_test_ModelWithExtensions* clone =
      (upb_test_ModelWithExtensions*)upb_Message_DeepClone(
          UPB_UPCAST(msg), &upb_0test__ModelWithExtensions_msg_init, arena);
  ASSERT_NE(clone, nullptr);

  // Mutate original
  upb_test_ModelExtension1_set_str(ext1, upb_StringView_FromString("Mutated"));
  upb_Arena_Free(source_arena);

  // Check if clone has the non-canonical extension and it's unmodified
  upb_MessageUnknown data;
  uintptr_t iter = kUpb_Message_UnknownBegin;
  bool has_non_canonical = false;
  const upb_Extension* ext_found = nullptr;
  while (upb_Message_NextUnknown2(UPB_UPCAST(clone), &data, &iter)) {
    if (data.type == kUpb_MessageUnknownType_NonCanonicalExtension) {
      has_non_canonical = true;
      ext_found = (const upb_Extension*)data.value.extension;
    }
  }
  EXPECT_TRUE(has_non_canonical);
  ASSERT_NE(ext_found, nullptr);

  const upb_test_ModelExtension1* cloned_ext =
      (const upb_test_ModelExtension1*)ext_found->data.msg_val;
  EXPECT_TRUE(
      upb_StringView_IsEqual(upb_test_ModelExtension1_str(cloned_ext),
                             upb_StringView_FromString("LifecycleValue")));

  upb_Arena_Free(arena);
}

}  // namespace

#include "upb/port/undef.inc"