python/free_threading/weak_map.h

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

#ifndef PYUPB_FREE_THREADING_WEAK_MAP_H_
#define PYUPB_FREE_THREADING_WEAK_MAP_H_

// clang-format off
#include "Python.h"
// clang-format on

#include <stdbool.h>
#include <stdint.h>

#include "python/free_threading/mutex.h"
#include "upb/hash/int_table.h"

// A WeakMap maps C pointers to the corresponding Python wrapper object. We
// want a consistent Python wrapper object for each C object, both to save
// memory and to provide object stability (ie. x is x).
//
// Each wrapped object should add itself to the map when it is constructed and
// remove itself from the map when it is destroyed. The map is weak so it does
// not take references to the cached objects.

struct PyUpb_WeakMap;
typedef struct PyUpb_WeakMap PyUpb_WeakMap;

PyUpb_WeakMap* PyUpb_WeakMap_New(void);
// Accepts NULL for cleanup of partially initialized objects.
void PyUpb_WeakMap_Free(PyUpb_WeakMap* map);

// Returns the number of entries in the map.
size_t PyUpb_WeakMap_Count(PyUpb_WeakMap* map);

// Adds the given object to the map, indexed by the given key.
// If key is already present in the map, py_obj is DECREF'd and replaced with
// the existing object.
//
// If memory allocation fails, returns false and does not modify the map.
bool PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject** obj);

// Adds the given object to the map, indexed by the given key.
// The key must not already be present in the map.
//
// If memory allocation fails, returns false and does not modify the map.
bool PyUpb_WeakMap_UniqueAdd(PyUpb_WeakMap* map, const void* key,
                             PyObject* obj);

// Removes the given key from the cache. It must exist in the cache currently.
void PyUpb_WeakMap_Erase(PyUpb_WeakMap* map, const void* key, PyObject* obj);

// Removes the entry from the cache, but only if it matches the given value.
// Returns true if the entry was found and removed.
bool PyUpb_WeakMap_EraseIfEqual(PyUpb_WeakMap* map, const void* key,
                                PyObject* obj);

// Returns a new reference to an object if it exists, otherwise returns NULL.
PyObject* PyUpb_WeakMap_Get(PyUpb_WeakMap* map, const void* key);

struct PyUpb_WeakMapEntry;
typedef struct PyUpb_WeakMapEntry PyUpb_WeakMapEntry;

typedef struct {
  PyUpb_WeakMap* map;
  PyUpb_WeakMapEntry* entries;
  size_t count;
  size_t index;
} PyUpb_WeakMapIter;

// Begin snapshots live entries with strong references under map->mutex, then
// releases the mutex before returning. The caller must retain the map owner
// until End. Next decrefs the previously yielded entry's snapshot reference,
// and End decrefs any remaining snapshot objects without holding a cache mutex.
// Delete only erases the current entry if the cache still contains that same
// object. Returns false with MemoryError on allocation failure, with no mutex
// held.
bool PyUpb_WeakMapIter_Begin(PyUpb_WeakMapIter* iter, PyUpb_WeakMap* map);
bool PyUpb_WeakMapIter_Next(PyUpb_WeakMapIter* iter, const void** key,
                            PyObject** obj);
void PyUpb_WeakMapIter_Delete(PyUpb_WeakMapIter* iter);
void PyUpb_WeakMapIter_End(PyUpb_WeakMapIter* iter);

#endif  // PYUPB_FREE_THREADING_WEAK_MAP_H_