upb/message/internal/extension.c

// 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

#include "upb/message/internal/extension.h"

#include <stdint.h>
#include <string.h>

#include "upb/mem/arena.h"
#include "upb/message/internal/extension.h"
#include "upb/message/internal/message.h"
#include "upb/message/internal/types.h"
#include "upb/mini_table/extension.h"

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

const upb_Extension* UPB_PRIVATE(_upb_Message_Getext)(
    const struct upb_Message* msg, const upb_MiniTableExtension* e) {
  upb_Message_Internal* in = UPB_PRIVATE(_upb_Message_GetInternal)(msg);
  if (!in) return NULL;

  for (size_t i = 0; i < in->size; i++) {
    upb_TaggedAuxPtr tagged_ptr = in->aux_data[i];
    if (upb_TaggedAuxPtr_IsCanonicalExtension(tagged_ptr)) {
      const upb_Extension* ext =
          upb_TaggedAuxPtr_CanonicalExtension(tagged_ptr);
      if (ext->ext == e) {
        return ext;
      }
    }
  }

  return NULL;
}

upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtensionWithTag)(
    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a,
    upb_TaggedAuxType tag) {
  UPB_ASSERT(!upb_Message_IsFrozen(msg));
  // For Canonical Extensions, we check whether the extension has already been
  // set. If we find an extension with the same pointer and tag, we reuse it to
  // prevent duplicate entries for the same extension.
  //
  // For Non-Canonical Extensions, we do NOT reuse them, matching the behavior
  // of adding a unknown StringView (through `_upb_Message_AddUnknown`) which
  // accumulates.
  if (tag == kUpb_TaggedAuxType_CanonicalExtension) {
    upb_Extension* ext =
        (upb_Extension*)UPB_PRIVATE(_upb_Message_Getext)(msg, e);
    if (ext) return ext;
  }
  if (!UPB_PRIVATE(_upb_Message_ReserveSlot)(msg, a)) return NULL;
  upb_Message_Internal* in = UPB_PRIVATE(_upb_Message_GetInternal)(msg);
  upb_Extension* ext = upb_Arena_Malloc(a, sizeof(upb_Extension));
  if (!ext) return NULL;
  memset(ext, 0, sizeof(upb_Extension));
  ext->ext = e;
  in->aux_data[in->size++] = upb_TaggedAuxPtr_MakeExtension(ext, tag);
  return ext;
}

upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a) {
  return UPB_PRIVATE(_upb_Message_GetOrCreateExtensionWithTag)(
      msg, e, a, kUpb_TaggedAuxType_CanonicalExtension);
}

upb_Extension* UPB_PRIVATE(_upb_Message_CreateNonCanonicalExtension)(
    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a) {
  return UPB_PRIVATE(_upb_Message_GetOrCreateExtensionWithTag)(
      msg, e, a, kUpb_TaggedAuxType_NonCanonicalExtension);
}