Significant memory-management refactoring any Python extension.

1 files changed, 92 insertions, 72 deletions

diff --git a/lang_ext/python/definition.c b/lang_ext/python/definition.c
index 6a788ab..cc1089d 100644
--- a/lang_ext/python/definition.c
+++ b/lang_ext/python/definition.c
@@ -24,21 +24,29 @@
 #include "upb_context.h"
 #include "upb_msg.h"
 
-#if PY_MAJOR_VERSION > 3
-const char *bytes_format = "y#";
-#else
-const char *bytes_format = "s#";
-#endif
+static PyTypeObject PyUpb_ContextType;
+static struct upb_strtable msgdefs;
+static struct upb_strtable contexts;
 
+struct msgtab_entry {
+  struct upb_strtable_entry e;
+  PyUpb_MsgDef *msgdef;
+};
 
-/* upb.def.MessageDefinition **************************************************/
+struct contexttab_entry {
+  struct upb_strtable_entry e;
+  PyUpb_Context *context;
+};
 
-typedef struct {
-  PyObject_HEAD
-  struct upb_msgdef *def;
-} PyUpb_MsgDef;
+#define CheckContext(obj) \
+  (void*)obj; do { \
+    if(!PyObject_TypeCheck(obj, &PyUpb_ContextType)) { \
+      PyErr_SetString(PyExc_TypeError, "Must be a upb.Context"); \
+      return NULL; \
+    } \
+  } while(0)
 
-PyTypeObject PyUpb_MsgDefType;  /* forward decl. */
+/* upb.def.MessageDefinition **************************************************/
 
 /* Not implemented yet, but these methods will expose information about the
  * message definition (the upb_msgdef). */
@@ -46,18 +54,10 @@ static PyMethodDef msgdef_methods[] = {
   {NULL, NULL}
 };
 
-static PyObject *msgdef_new(struct upb_msgdef *m)
-{
-  PyUpb_MsgDef *md_obj = (void*)PyType_GenericAlloc(&PyUpb_MsgDefType, 0);
-  md_obj->def = m;
-  upb_msgdef_ref(md_obj->def);
-  return (void*)md_obj;
-}
-
 static void msgdef_dealloc(PyObject *obj)
 {
   PyUpb_MsgDef *md_obj = (void*)obj;
-  upb_msgdef_unref(md_obj->def);
+  Py_DECREF(md_obj->context);
   obj->ob_type->tp_free(obj);
 }
 
@@ -106,27 +106,11 @@ PyTypeObject PyUpb_MsgDefType = {
 
 /* upb.Context ****************************************************************/
 
-typedef struct {
-  PyObject_HEAD
-  struct upb_context *context;
-  PyObject *created_defs;
-} PyUpb_Context;
-
-static PyTypeObject PyUpb_ContextType;  /* forward decl. */
-
-#define CheckContext(obj) \
-  (void*)obj; do { \
-    if(!PyObject_TypeCheck(obj, &PyUpb_ContextType)) { \
-      PyErr_SetString(PyExc_TypeError, "Must be a upb.Context"); \
-      return NULL; \
-    } \
-  } while(0)
-
 static PyObject *context_parsefds(PyObject *_context, PyObject *args)
 {
   PyUpb_Context *context = CheckContext(_context);
   struct upb_string str;
-  if(!PyArg_ParseTuple(args, bytes_format, &str.ptr, &str.byte_len))
+  if(!PyArg_ParseTuple(args, BYTES_FORMAT, &str.ptr, &str.byte_len))
     return NULL;
   str.byte_size = 0;  /* We don't own that mem. */
 
@@ -138,35 +122,56 @@ static PyObject *context_parsefds(PyObject *_context, PyObject *args)
   Py_RETURN_NONE;
 }
 
-static PyObject *get_or_create_def(PyUpb_Context *context,
-                                   struct upb_symtab_entry *e)
+static PyObject *get_or_create_def(struct upb_symtab_entry *e)
 {
-  /* Check out internal dictionary of Python classes we have already created
-   * (keyed by the address of the obj we are referencing). */
-#if PY_MAJOR_VERSION > 3
-  PyObject *str = PyBytes_FromStringAndSize((char*)&e->ref, sizeof(void*));
-#else
-  PyObject *str = PyString_FromStringAndSize((char*)&e->ref, sizeof(void*));
-#endif
-  /* Would use PyDict_GetItemStringAndSize() if it existed, but only
-   * PyDict_GetItemString() exists, and pointers could have NULL bytes. */
-  PyObject *def = PyDict_GetItem(context->created_defs, str);
-  if(!def) {
-    switch(e->type) {
-      case UPB_SYM_MESSAGE:
-        def = msgdef_new(e->ref.msg);
-        break;
-      case UPB_SYM_ENUM:
-      case UPB_SYM_SERVICE:
-      case UPB_SYM_EXTENSION:
-      default:
-        def = NULL;
-        break;
-    }
-    if(def) PyDict_SetItem(context->created_defs, str, def);
+  switch(e->type) {
+    case UPB_SYM_MESSAGE: return (PyObject*)get_or_create_msgdef(e->ref.msg);
+    case UPB_SYM_ENUM:
+    case UPB_SYM_SERVICE:
+    case UPB_SYM_EXTENSION:
+    default: fprintf(stderr, "upb.pb, not implemented.\n"); abort(); return NULL;
   }
-  Py_DECREF(str);
-  return def;
+}
+
+static PyUpb_Context *get_or_create_context(struct upb_context *context)
+{
+  PyUpb_Context *pycontext = NULL;
+  struct upb_string str = {.ptr = (char*)&context, .byte_len = sizeof(void*)};
+  struct contexttab_entry *e = upb_strtable_lookup(&contexts, &str);
+  if(!e) {
+    pycontext = (void*)PyUpb_ContextType.tp_alloc(&PyUpb_ContextType, 0);
+    pycontext->context = context;
+    struct contexttab_entry new_e = {
+      .e = {.key = {.ptr = (char*)&pycontext->context, .byte_len = sizeof(void*)}},
+      .context = pycontext
+    };
+    upb_strtable_insert(&contexts, &new_e.e);
+  } else {
+    pycontext = e->context;
+    Py_INCREF(pycontext);
+  }
+  return pycontext;
+}
+
+PyUpb_MsgDef *get_or_create_msgdef(struct upb_msgdef *def)
+{
+  PyUpb_MsgDef *pydef = NULL;
+  struct upb_string str = {.ptr = (char*)&def, .byte_len = sizeof(void*)};
+  struct msgtab_entry *e = upb_strtable_lookup(&msgdefs, &str);
+  if(!e) {
+    pydef = (void*)PyUpb_MsgDefType.tp_alloc(&PyUpb_MsgDefType, 0);
+    pydef->def = def;
+    pydef->context = get_or_create_context(def->context);
+    struct msgtab_entry new_e = {
+      .e = {.key = {.ptr = (char*)&pydef->def, .byte_len = sizeof(void*)}},
+      .msgdef = pydef
+    };
+    upb_strtable_insert(&msgdefs, &new_e.e);
+  } else {
+    pydef = e->msgdef;
+    Py_INCREF(pydef);
+  }
+  return pydef;
 }
 
 static PyObject *context_lookup(PyObject *self, PyObject *args)
@@ -179,7 +184,7 @@ static PyObject *context_lookup(PyObject *self, PyObject *args)
 
   struct upb_symtab_entry e;
   if(upb_context_lookup(context->context, &str, &e)) {
-    return get_or_create_def(context, &e);
+    return get_or_create_def(&e);
   } else {
     Py_RETURN_NONE;
   }
@@ -197,12 +202,13 @@ static PyObject *context_resolve(PyObject *self, PyObject *args)
 
   struct upb_symtab_entry e;
   if(upb_context_resolve(context->context, &base, &str, &e)) {
-    return get_or_create_def(context, &e);
+    return get_or_create_def(&e);
   } else {
     Py_RETURN_NONE;
   }
 }
 
+/* Callback for upb_context_enumerate below. */
 static void add_string(void *udata, struct upb_symtab_entry *entry)
 {
   PyObject *list = udata;
@@ -244,7 +250,11 @@ static PyObject *context_new(PyTypeObject *subtype,
 {
   PyUpb_Context *obj = (void*)subtype->tp_alloc(subtype, 0);
   obj->context = upb_context_new();
-  obj->created_defs = PyDict_New();
+  struct contexttab_entry e = {
+    .e = {.key = {.ptr = (char*)&obj->context, .byte_len = sizeof(void*)}},
+    .context = obj
+  };
+  upb_strtable_insert(&contexts, &e.e);
   return (void*)obj;
 }
 
@@ -252,7 +262,9 @@ static void context_dealloc(PyObject *obj)
 {
   PyUpb_Context *c = (void*)obj;
   upb_context_unref(c->context);
-  Py_DECREF(c->created_defs);
+  /* TODO: once strtable supports delete. */
+  //struct upb_string ptrstr = {.ptr = (char*)&c->context, .byte_len = sizeof(void*)};
+  //upb_strtable_delete(&contexts, &ptrstr);
   obj->ob_type->tp_free(obj);
 }
 
@@ -299,17 +311,25 @@ static PyTypeObject PyUpb_ContextType = {
   0,                                      /* tp_free */
 };
 
-PyMethodDef methods[] = {
+static PyMethodDef methods[] = {
+  {NULL, NULL}
 };
 
 PyMODINIT_FUNC
 initdefinition(void)
 {
   if(PyType_Ready(&PyUpb_ContextType) < 0) return;
-  Py_INCREF(&PyUpb_ContextType);  /* TODO: necessary? */
   if(PyType_Ready(&PyUpb_MsgDefType) < 0) return;
-  Py_INCREF(&PyUpb_MsgDefType);  /* TODO: necessary? */
 
-  PyObject *mod = Py_InitModule("upb.definition", methods);
+  /* PyModule_AddObject steals a reference.  These objects are statically
+   * allocated and must not be deleted, so we increment their refcount. */
+  Py_INCREF(&PyUpb_ContextType);
+  Py_INCREF(&PyUpb_MsgDefType);
+
+  PyObject *mod = Py_InitModule("upb.cext.definition", methods);
   PyModule_AddObject(mod, "Context", (PyObject*)&PyUpb_ContextType);
+  PyModule_AddObject(mod, "MessageDefinition", (PyObject*)&PyUpb_MsgDefType);
+
+  upb_strtable_init(&contexts, 8, sizeof(struct contexttab_entry));
+  upb_strtable_init(&msgdefs, 16, sizeof(struct msgtab_entry));
 }