upb/mini_table/extension_registry.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/mini_table/extension_registry.h"

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

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

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

struct upb_ExtensionRegistry {
  upb_exttable exts;
  upb_Arena* arena;
};

upb_ExtensionRegistry* upb_ExtensionRegistry_New(upb_Arena* arena) {
  upb_ExtensionRegistry* r = upb_Arena_Malloc(arena, sizeof(*r));
  if (!r) return NULL;
  r->arena = arena;
  if (!upb_exttable_init(&r->exts, 8, arena)) return NULL;
  return r;
}

UPB_API upb_ExtensionRegistryStatus upb_ExtensionRegistry_Add(
    upb_ExtensionRegistry* r, const upb_MiniTableExtension* e) {
  UPB_STATIC_ASSERT(
      offsetof(upb_MiniTableExtension,
               UPB_PRIVATE(field).UPB_PRIVATE(number)) == 0,
      "Extension must be first-member-of-struct convertable with uint32_t");
  uint32_t fieldnum = upb_MiniTableExtension_Number(e);
  const upb_MiniTable* extendee = upb_MiniTableExtension_Extendee(e);

  const uint32_t kMaxFieldNumber = (1 << 29) - 1;
  if (fieldnum == 0 ||
      (fieldnum > kMaxFieldNumber && !upb_MiniTable_IsMessageSet(extendee))) {
    return kUpb_ExtensionRegistryStatus_InvalidExtension;
  }

  UPB_ASSERT(upb_MiniTable_FindFieldByNumber(extendee, fieldnum) == NULL);

  if (upb_exttable_lookup(&r->exts, extendee, fieldnum) != NULL) {
    return kUpb_ExtensionRegistryStatus_DuplicateEntry;
  }

  if (!upb_exttable_insert(&r->exts, extendee, (const uint32_t*)e, r->arena)) {
    return kUpb_ExtensionRegistryStatus_OutOfMemory;
  }
  return kUpb_ExtensionRegistryStatus_Ok;
}

upb_ExtensionRegistryStatus upb_ExtensionRegistry_AddArray(
    upb_ExtensionRegistry* r, const upb_MiniTableExtension** e, size_t count) {
  const upb_MiniTableExtension** start = e;
  const upb_MiniTableExtension** end = UPB_PTRADD(e, count);
  upb_ExtensionRegistryStatus status = kUpb_ExtensionRegistryStatus_Ok;
  for (; e < end; e++) {
    status = upb_ExtensionRegistry_Add(r, *e);
    if (status != kUpb_ExtensionRegistryStatus_Ok) goto failure;
  }
  return kUpb_ExtensionRegistryStatus_Ok;

failure:
  // Back out the entries previously added.
  for (end = e, e = start; e < end; e++) {
    const upb_MiniTableExtension* ext = *e;
    upb_exttable_remove(&r->exts, upb_MiniTableExtension_Extendee(ext),
                        upb_MiniTableExtension_Number(ext));
  }
  UPB_ASSERT(status != kUpb_ExtensionRegistryStatus_Ok);
  return status;
}

const upb_MiniTableExtension* upb_ExtensionRegistry_Lookup(
    const upb_ExtensionRegistry* r, const upb_MiniTable* t, uint32_t num) {
  const uint32_t* v = upb_exttable_lookup(&r->exts, t, num);
  return (const upb_MiniTableExtension*)v;
}

size_t upb_ExtensionRegistry_Size(const upb_ExtensionRegistry* r) {
  return upb_exttable_size(&r->exts);
}