// 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 "python/map.h"

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

#include "google/protobuf/breaking_changes.h"
#include "python/convert.h"
#include "python/descriptor.h"
#include "python/free_threading/weak_map.h"
#include "python/message.h"
#include "python/protobuf.h"
#include "upb/mem/arena.h"
#include "upb/message/array.h"
#include "upb/message/map.h"
#include "upb/message/message.h"
#include "upb/mini_table/message.h"
#include "upb/reflection/def.h"

// -----------------------------------------------------------------------------
// MapContainer
// -----------------------------------------------------------------------------

typedef struct {
  // clang-format off
  PyObject_HEAD
  PyObject* arena;
  // clang-format on
  // The field descriptor (upb_FieldDef*).
  // The low bit indicates whether the container is reified (see ptr below).
  //   - low bit set: repeated field is a stub (empty map, no underlying data).
  //   - low bit clear: repeated field is reified (points to upb_Array).
  uintptr_t field;
  union {
    PyObject* parent;  // stub: owning pointer to parent message.
    upb_Map* map;      // reified: the data for this array.
  } ptr;
  int version;
} PyUpb_MapContainer;

static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);

static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
  return self->field & 1;
}

static PyObject* PyUpb_MapContainer_GetPool(PyUpb_MapContainer* self) {
  return PyUpb_Arena_GetPool(self->arena);
}

// If the map is reified, returns it.  Otherwise, returns NULL.
// If NULL is returned, the object is empty and has no underlying data.
static upb_Map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
  return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
}

static const upb_FieldDef* PyUpb_MapContainer_GetField(
    PyUpb_MapContainer* self) {
  return (const upb_FieldDef*)(self->field & ~(uintptr_t)1);
}

static void PyUpb_MapContainer_Dealloc(void* _self) {
  PyUpb_MapContainer* self = _self;
  if (PyUpb_MapContainer_IsStub(self)) {
    PyUpb_Message_CacheDelete(self->ptr.parent,
                              PyUpb_MapContainer_GetField(self), _self);
    Py_DECREF(self->ptr.parent);
  } else {
    PyUpb_Arena_CacheEraseIfEqual(self->arena, self->ptr.map, _self);
  }
  Py_DECREF(self->arena);
  PyUpb_Dealloc(_self);
}

static PyTypeObject* PyUpb_MapContainer_GetClass(const upb_FieldDef* f) {
  assert(upb_FieldDef_IsMap(f));
  PyUpb_ModuleState* state = PyUpb_ModuleState_MaybeGet();
  if (!state) return NULL;
  const upb_FieldDef* val =
      upb_MessageDef_Field(upb_FieldDef_MessageSubDef(f), 1);
  assert(upb_FieldDef_Number(val) == 2);
  return upb_FieldDef_IsSubMessage(val) ? state->message_map_container_type
                                        : state->scalar_map_container_type;
}

PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
                                     PyObject* arena) {
  PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
  if (!cls) {
    PyErr_SetString(PyExc_RuntimeError, "Interpreter is finalizing");
    return NULL;
  }
  PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
  if (map == NULL) return NULL;
  map->arena = arena;
  map->field = (uintptr_t)f | 1;
  map->ptr.parent = parent;
  map->version = 0;
  Py_INCREF(arena);
  Py_INCREF(parent);
  return &map->ob_base;
}

/*
 * PyUpb_MapContainer_Reify()
 *
 * Transitions the wrapper from the unset "stub" state (owning a reference on
 * self->ptr.parent) to the set state (having a non-owning pointer to
 * self->ptr.map).
 *
 * If `map` is NULL, a new map is allocated, otherwise the wrapper will be
 * attached to the provided `arr`.
 *
 * If `subobj_map` and `iter` are non-NULL, the wrapper is assumed to be in the
 * unset state, and the entry for `iter` is removed from `subobj_map`.
 */
upb_Map* PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map,
                                  PyUpb_WeakMap* subobj_map, intptr_t* iter) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  assert(PyUpb_MapContainer_IsStub(self));
  if (!map) {
    const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
    upb_Arena* arena = PyUpb_Arena_Get(self->arena);
    const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
    const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
    const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
    map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
                      upb_FieldDef_CType(val_f));
    if (!map) {
      PyErr_SetNone(PyExc_MemoryError);
      return NULL;
    }
  }
  if (subobj_map) {
    PyUpb_WeakMap_DeleteIter(subobj_map, iter);
  } else {
    const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
    upb_MessageValue msgval = {.map_val = map};
    if (!PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval, _self)) {
      return NULL;
    }
  }
  if (!PyUpb_Arena_CacheUniqueAdd(self->arena, map, &self->ob_base)) {
    return NULL;
  }
  Py_DECREF(self->ptr.parent);
  self->ptr.map = map;  // Overwrites self->ptr.parent.
  self->field &= ~(uintptr_t)1;
  assert(!PyUpb_MapContainer_IsStub(self));
  return map;
}

void PyUpb_MapContainer_Invalidate(PyObject* obj) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
  self->version++;
}

bool PyUpb_MapContainer_IsFrozen(PyUpb_MapContainer* self) {
  if (PyUpb_MapContainer_IsStub(self)) {
    return PyUpb_Message_IsFrozen(self->ptr.parent);
  } else {
    return upb_Map_IsFrozen(self->ptr.map) || PyUpb_Arena_IsFrozen(self->arena);
  }
}

upb_Map* PyUpb_MapContainer_AssureWritable(PyObject* _self) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  if (PyUpb_MapContainer_IsFrozen(self)) {
    if (!PyUpb_CheckFrozen(true, "Map is immutable")) {
      return NULL;
    }
  }

  self->version++;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  if (map) return map;  // Already writable.

  return PyUpb_MapContainer_Reify((PyObject*)self, NULL, NULL, 0);
}

static bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,
                                   upb_MessageValue key, upb_MessageValue val,
                                   upb_Arena* arena) {
  switch (upb_Map_Insert(map, key, val, arena)) {
    case kUpb_MapInsertStatus_Inserted:
      return true;
    case kUpb_MapInsertStatus_Replaced:
      // We did not insert a new key, undo the previous invalidate.
      self->version--;
      return true;
    case kUpb_MapInsertStatus_OutOfMemory:
      PyErr_SetNone(PyExc_MemoryError);
      return false;
  }
  return false;  // Unreachable, silence compiler warning.
}

// Assigns `self[key] = val` for the map `self`.
static int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
                                              PyObject* val) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_AssureWritable(_self);
  if (!map) return -1;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
  upb_Arena* arena = PyUpb_Arena_Get(self->arena);
  upb_MessageValue u_key, u_val;
  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return -1;

  if (val) {
    if (!PyUpb_PyToUpb(val, val_f, &u_val, arena)) return -1;
    if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return -1;
  } else {
    if (!upb_Map_Delete(map, u_key, NULL)) {
      PyErr_Format(PyExc_KeyError, "Key not present in map");
      return -1;
    }
  }
  return 0;
}

static PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
  upb_MessageValue u_key, u_val;
  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
  if (!map || !upb_Map_Get(map, u_key, &u_val)) {
    map = PyUpb_MapContainer_AssureWritable(_self);
    if (!map) return NULL;
    upb_Arena* arena = PyUpb_Arena_Get(self->arena);
    if (upb_FieldDef_IsSubMessage(val_f)) {
      const upb_MessageDef* m = upb_FieldDef_MessageSubDef(val_f);
      const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
      u_val.msg_val = upb_Message_New(layout, arena);
      if (!u_val.msg_val) {
        PyErr_SetNone(PyExc_MemoryError);
        return NULL;
      }
    } else {
      memset(&u_val, 0, sizeof(u_val));
    }
    if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return NULL;
  }
  return PyUpb_UpbToPy(u_val, val_f, self->arena);
}

static int PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  if (!map) return 0;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  upb_MessageValue u_key;
  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return -1;
  if (upb_Map_Get(map, u_key, NULL)) {
    return 1;
  } else {
    return 0;
  }
}

static PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
  upb_Map* map = PyUpb_MapContainer_AssureWritable(_self);
  if (!map) return NULL;
  // TODO: b/517235198 - Reify even for empty sequences.
  if (upb_Map_Size(map) > 0) {
    upb_Map_Clear(map);
  }
  Py_RETURN_NONE;
}

static PyObject* PyUpb_ScalarMapContainer_Setdefault(PyObject* _self,
                                                     PyObject* args) {
  PyObject* key;
  PyObject* default_value = Py_None;

  if (!PyArg_UnpackTuple(args, "setdefault", 1, 2, &key, &default_value)) {
    return NULL;
  }

  if (default_value == Py_None) {
    PyErr_Format(PyExc_ValueError,
                 "The value for scalar map setdefault must be set.");
    return NULL;
  }

  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_AssureWritable(_self);
  if (!map) return NULL;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
  upb_MessageValue u_key, u_val;
  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
  if (map && upb_Map_Get(map, u_key, &u_val)) {
    return PyUpb_UpbToPy(u_val, val_f, self->arena);
  }

  upb_Arena* arena = PyUpb_Arena_Get(self->arena);
  if (!PyUpb_PyToUpb(default_value, val_f, &u_val, arena)) return NULL;
  if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return NULL;

  Py_INCREF(default_value);
  return default_value;
}

static PyObject* PyUpb_MessageMapContainer_Setdefault(PyObject* self,
                                                      PyObject* args) {
  PyErr_Format(PyExc_NotImplementedError,
               "Set message map value directly is not supported.");
  return NULL;
}

static PyObject* PyUpb_MapContainer_Get(PyObject* _self, PyObject* args,
                                        PyObject* kwargs) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  static const char* kwlist[] = {"key", "default", NULL};
  PyObject* key;
  PyObject* default_value = NULL;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
                                   &default_value)) {
    return NULL;
  }

  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
  upb_MessageValue u_key, u_val;
  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
  if (map && upb_Map_Get(map, u_key, &u_val)) {
    return PyUpb_UpbToPy(u_val, val_f, self->arena);
  }
  if (default_value) {
    Py_INCREF(default_value);
    return default_value;
  }
  Py_RETURN_NONE;
}

static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
                                                  PyObject* arg) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  PyObject* pool = PyUpb_MapContainer_GetPool(self);
  return PyUpb_Descriptor_GetClass(pool, entry_m);
}

static Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  return map ? upb_Map_Size(map) : 0;
}

int PyUpb_Message_InitMapAttributes(PyObject* map, PyObject* value,
                                    const upb_FieldDef* f);

static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_AssureWritable(_self);
  if (!map) return NULL;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);

  if (PyDict_Check(_arg)) {
    return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
  }

  if (PyUpb_Message_InitMapAttributes(_self, _arg, f) < 0) {
    return NULL;
  }

  Py_RETURN_NONE;
}

static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
  PyObject* dict = PyDict_New();
  if (map) {
    const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
    const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
    const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
    const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
    size_t iter = kUpb_Map_Begin;
    upb_MessageValue map_key, map_val;
    while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
      PyObject* key = PyUpb_UpbToPy(map_key, key_f, self->arena);
      PyObject* val = PyUpb_UpbToPy(map_val, val_f, self->arena);
      if (!key || !val) {
        Py_XDECREF(key);
        Py_XDECREF(val);
        Py_DECREF(dict);
        return NULL;
      }
      PyDict_SetItem(dict, key, val);
      Py_DECREF(key);
      Py_DECREF(val);
    }
  }
  PyObject* repr = PyObject_Repr(dict);
  Py_DECREF(dict);
  return repr;
}

PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
                                                const upb_FieldDef* f,
                                                PyObject* arena) {
  PyObject* ret = PyUpb_Arena_CacheGet(arena, map);
  if (ret) return ret;

  PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
  if (!cls) {
    PyErr_SetString(PyExc_RuntimeError, "Interpreter is finalizing");
    return NULL;
  }
  PyUpb_MapContainer* map_cont = (void*)PyType_GenericAlloc(cls, 0);
  if (map_cont == NULL) return NULL;
  map_cont->arena = arena;
  map_cont->field = (uintptr_t)f;
  map_cont->ptr.map = map;
  map_cont->version = 0;
  ret = &map_cont->ob_base;
  Py_INCREF(arena);
  if (!PyUpb_Arena_CacheAdd(arena, map, &ret)) {
    Py_DECREF(ret);
    return NULL;
  }
  return ret;
}

// -----------------------------------------------------------------------------
// ScalarMapContainer
// -----------------------------------------------------------------------------

static PyMethodDef PyUpb_ScalarMapContainer_Methods[] = {
    {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
     "Removes all elements from the map."},
    {"setdefault", (PyCFunction)PyUpb_ScalarMapContainer_Setdefault,
     METH_VARARGS,
     "If the key does not exist, insert the key, with the specified value"},
    {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
     "Gets the value for the given key if present, or otherwise a default"},
    {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
     "Return the class used to build Entries of (key, value) pairs."},
    {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
     "Merges a map into the current map."},
    /*
   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
     "Makes a deep copy of the class." },
   { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
     "Outputs picklable representation of the repeated field." },
   */
    {NULL, NULL},
};

static PyType_Slot PyUpb_ScalarMapContainer_Slots[] = {
    {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
    {Py_mp_length, PyUpb_MapContainer_Length},
    {Py_mp_subscript, PyUpb_MapContainer_Subscript},
    {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
    {Py_sq_contains, PyUpb_MapContainer_Contains},
    {Py_tp_methods, PyUpb_ScalarMapContainer_Methods},
    {Py_tp_iter, PyUpb_MapIterator_New},
    {Py_tp_repr, PyUpb_MapContainer_Repr},
    {0, NULL},
};

static PyType_Spec PyUpb_ScalarMapContainer_Spec = {
    PYUPB_MODULE_NAME ".ScalarMapContainer",
    sizeof(PyUpb_MapContainer),
    0,
    Py_TPFLAGS_DEFAULT,
    PyUpb_ScalarMapContainer_Slots,
};

// -----------------------------------------------------------------------------
// MessageMapContainer
// -----------------------------------------------------------------------------

static PyMethodDef PyUpb_MessageMapContainer_Methods[] = {
    {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
     "Removes all elements from the map."},
    {"setdefault", (PyCFunction)PyUpb_MessageMapContainer_Setdefault,
     METH_VARARGS, "setdefault is disallowed in MessageMap."},
    {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
     "Gets the value for the given key if present, or otherwise a default"},
    {"get_or_create", PyUpb_MapContainer_Subscript, METH_O,
     "Alias for getitem, useful to make explicit that the map is mutated."},
    {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
     "Return the class used to build Entries of (key, value) pairs."},
    {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
     "Merges a map into the current map."},
    /*
   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
     "Makes a deep copy of the class." },
   { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
     "Outputs picklable representation of the repeated field." },
   */
    {NULL, NULL},
};

static PyType_Slot PyUpb_MessageMapContainer_Slots[] = {
    {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
    {Py_mp_length, PyUpb_MapContainer_Length},
    {Py_mp_subscript, PyUpb_MapContainer_Subscript},
    {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
    {Py_sq_contains, PyUpb_MapContainer_Contains},
    {Py_tp_methods, PyUpb_MessageMapContainer_Methods},
    {Py_tp_iter, PyUpb_MapIterator_New},
    {Py_tp_repr, PyUpb_MapContainer_Repr},
    {0, NULL}};

static PyType_Spec PyUpb_MessageMapContainer_Spec = {
    PYUPB_MODULE_NAME ".MessageMapContainer", sizeof(PyUpb_MapContainer), 0,
    Py_TPFLAGS_DEFAULT, PyUpb_MessageMapContainer_Slots};

// -----------------------------------------------------------------------------
// MapIterator
// -----------------------------------------------------------------------------

typedef struct {
  // clang-format off
  PyObject_HEAD
  PyUpb_MapContainer* map;  // We own a reference.
  // clang-format on
  size_t iter;
  int version;
} PyUpb_MapIterator;

static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map) {
  PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
  PyUpb_MapIterator* iter =
      (void*)PyType_GenericAlloc(state->map_iterator_type, 0);
  iter->map = map;
  iter->iter = kUpb_Map_Begin;
  iter->version = map->version;
  Py_INCREF(map);
  return &iter->ob_base;
}

static void PyUpb_MapIterator_Dealloc(void* _self) {
  PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
  Py_DECREF(&self->map->ob_base);
  PyUpb_Dealloc(_self);
}

static PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
  PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
  if (self->version != self->map->version) {
    return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration.");
  }
  upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
  if (!map) return NULL;
  upb_MessageValue key, val;
  if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self->map);
  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
  return PyUpb_UpbToPy(key, key_f, self->map->arena);
}

static PyType_Slot PyUpb_MapIterator_Slots[] = {
    {Py_tp_dealloc, PyUpb_MapIterator_Dealloc},
    {Py_tp_iter, PyObject_SelfIter},
    {Py_tp_iternext, PyUpb_MapIterator_IterNext},
    {0, NULL}};

static PyType_Spec PyUpb_MapIterator_Spec = {
    PYUPB_MODULE_NAME ".MapIterator", sizeof(PyUpb_MapIterator), 0,
    Py_TPFLAGS_DEFAULT, PyUpb_MapIterator_Slots};

// -----------------------------------------------------------------------------
// Top Level
// -----------------------------------------------------------------------------

static PyObject* GetMutableMappingBase(void) {
  PyObject* collections = NULL;
  PyObject* mapping = NULL;
  PyObject* base = NULL;
  if ((collections = PyImport_ImportModule("collections.abc")) &&
      (mapping = PyObject_GetAttrString(collections, "MutableMapping"))) {
    base = Py_BuildValue("O", mapping);
  }
  Py_XDECREF(collections);
  Py_XDECREF(mapping);
  return base;
}

bool PyUpb_Map_Init(PyObject* m) {
  PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
  PyObject* base = GetMutableMappingBase();
  if (!base) return false;

  const char* methods[] = {"keys", "items",   "values", "__eq__", "__ne__",
                           "pop",  "popitem", "update", NULL};

  state->message_map_container_type = PyUpb_AddClassWithRegister(
      m, &PyUpb_MessageMapContainer_Spec, base, methods);
  if (!state->message_map_container_type) {
    return false;
  }
  state->scalar_map_container_type = PyUpb_AddClassWithRegister(
      m, &PyUpb_ScalarMapContainer_Spec, base, methods);
  if (!state->scalar_map_container_type) {
    return false;
  }
  state->map_iterator_type = PyUpb_AddClass(m, &PyUpb_MapIterator_Spec);

  Py_DECREF(base);

  return state->message_map_container_type &&
         state->scalar_map_container_type && state->map_iterator_type;
}