Some refcounting fixes.

Clearly this stuff is too complex overall.
The plan is to move away from this and more towards
pools, like proto2 uses.

4 files changed, 32 insertions, 8 deletions

diff --git a/tests/test_def.c b/tests/test_def.c
index e9ca438..669aed0 100644
--- a/tests/test_def.c
+++ b/tests/test_def.c
@@ -53,9 +53,11 @@ static upb_symtab *load_test_proto(void *owner) {
 
   files_ptr = files;
   while (*files_ptr) {
-    bool ok = upb_symtab_addfile(s, *files, &status);
+    ASSERT(!upb_filedef_isfrozen(*files_ptr));
+    bool ok = upb_symtab_addfile(s, *files_ptr, &status);
     ASSERT(ok);
-    upb_filedef_unref(*files, &files);
+    ASSERT(upb_filedef_isfrozen(*files_ptr));
+    upb_filedef_unref(*files_ptr, &files);
     files_ptr++;
   }
 
diff --git a/upb/def.c b/upb/def.c
index c47f020..92e89e0 100644
--- a/upb/def.c
+++ b/upb/def.c
@@ -439,7 +439,16 @@ bool upb_def_freeze(upb_def *const* defs, size_t n, upb_status *s) {
 
 /* upb_enumdef ****************************************************************/
 
-static void upb_enumdef_free(upb_refcounted *r) {
+static void visitenum(const upb_refcounted *r, upb_refcounted_visit *visit,
+                      void *closure) {
+  const upb_enumdef *e = (const upb_enumdef*)r;
+  const upb_def *def = upb_enumdef_upcast(e);
+  if (upb_def_file(def)) {
+    visit(r, upb_filedef_upcast(upb_def_file(def)), closure);
+  }
+}
+
+static void freeenum(upb_refcounted *r) {
   upb_enumdef *e = (upb_enumdef*)r;
   upb_inttable_iter i;
   upb_inttable_begin(&i, &e->iton);
@@ -453,7 +462,7 @@ static void upb_enumdef_free(upb_refcounted *r) {
   upb_gfree(e);
 }
 
-const struct upb_refcounted_vtbl upb_enumdef_vtbl = {NULL, &upb_enumdef_free};
+const struct upb_refcounted_vtbl upb_enumdef_vtbl = {&visitenum, &freeenum};
 
 upb_enumdef *upb_enumdef_new(const void *owner) {
   upb_enumdef *e = upb_gmalloc(sizeof(*e));
@@ -611,6 +620,7 @@ const char *upb_fielddef_fullname(const upb_fielddef *e) {
 static void visitfield(const upb_refcounted *r, upb_refcounted_visit *visit,
                        void *closure) {
   const upb_fielddef *f = (const upb_fielddef*)r;
+  const upb_def *def = upb_fielddef_upcast(f);
   if (upb_fielddef_containingtype(f)) {
     visit(r, upb_msgdef_upcast2(upb_fielddef_containingtype(f)), closure);
   }
@@ -620,6 +630,9 @@ static void visitfield(const upb_refcounted *r, upb_refcounted_visit *visit,
   if (upb_fielddef_subdef(f)) {
     visit(r, upb_def_upcast(upb_fielddef_subdef(f)), closure);
   }
+  if (upb_def_file(def)) {
+    visit(r, upb_filedef_upcast(upb_def_file(def)), closure);
+  }
 }
 
 static void freefield(upb_refcounted *r) {
@@ -1384,6 +1397,7 @@ static void visitmsg(const upb_refcounted *r, upb_refcounted_visit *visit,
                      void *closure) {
   upb_msg_oneof_iter o;
   const upb_msgdef *m = (const upb_msgdef*)r;
+  const upb_def *def = upb_msgdef_upcast(m);
   upb_msg_field_iter i;
   for(upb_msg_field_begin(&i, m);
       !upb_msg_field_done(&i);
@@ -1397,6 +1411,9 @@ static void visitmsg(const upb_refcounted *r, upb_refcounted_visit *visit,
     upb_oneofdef *f = upb_msg_iter_oneof(&o);
     visit(r, upb_oneofdef_upcast(f), closure);
   }
+  if (upb_def_file(def)) {
+    visit(r, upb_filedef_upcast(upb_def_file(def)), closure);
+  }
 }
 
 static void freemsg(upb_refcounted *r) {
@@ -1545,6 +1562,7 @@ bool upb_msgdef_addfield(upb_msgdef *m, upb_fielddef *f, const void *ref_donor,
    * This method is idempotent. Check if |f| is already part of this msgdef and
    * return immediately if so. */
   if (upb_fielddef_containingtype(f) == m) {
+    if (ref_donor) upb_fielddef_unref(f, ref_donor);
     return true;
   }
 
@@ -2089,6 +2107,7 @@ bool upb_filedef_adddef(upb_filedef *f, upb_def *def, const void *ref_donor,
   if (upb_inttable_push(&f->defs, upb_value_constptr(def))) {
     def->file = f;
     upb_ref2(def, f);
+    upb_ref2(f, def);
     if (ref_donor) upb_def_unref(def, ref_donor);
     if (def->type == UPB_DEF_MSG) {
       upb_downcast_msgdef_mutable(def)->syntax = f->syntax;
diff --git a/upb/descriptor/reader.c b/upb/descriptor/reader.c
index b515840..2b9e0e9 100644
--- a/upb/descriptor/reader.c
+++ b/upb/descriptor/reader.c
@@ -573,7 +573,7 @@ static size_t field_ondefaultval(void *closure, const void *hd, const char *buf,
 static bool field_ononeofindex(void *closure, const void *hd, int32_t index) {
   upb_descreader *r = closure;
   upb_oneofdef *o = upb_descreader_getoneof(r, index);
-  bool ok = upb_oneofdef_addfield(o, r->f, NULL, NULL);
+  bool ok = upb_oneofdef_addfield(o, r->f, &r->f, NULL);
   UPB_UNUSED(hd);
 
   UPB_ASSERT(ok);
@@ -651,6 +651,7 @@ static void *msg_startext(void *closure, const void *hd) {
   return r;
 }
 
+#include <stdio.h>
 static void *msg_startfield(void *closure, const void *hd) {
   upb_descreader *r = closure;
   r->f = upb_fielddef_new(&r->f);
@@ -663,9 +664,13 @@ static void *msg_startfield(void *closure, const void *hd) {
 static bool msg_endfield(void *closure, const void *hd) {
   upb_descreader *r = closure;
   upb_msgdef *m = upb_descreader_top(r);
+  bool ok;
   UPB_UNUSED(hd);
 
-  upb_msgdef_addfield(m, r->f, &r->f, NULL);
+  if (upb_fielddef_containingoneof(r->f) == NULL) {
+    ok = upb_msgdef_addfield(m, r->f, &r->f, NULL);
+    UPB_ASSERT(ok);
+  }
   r->f = NULL;
   return true;
 }
diff --git a/upb/refcounted.h b/upb/refcounted.h
index 6698d38..1537414 100644
--- a/upb/refcounted.h
+++ b/upb/refcounted.h
@@ -28,8 +28,6 @@
  * For this reason we don't enable it by default, even in debug builds.
  */
 
-/* #define UPB_DEBUG_REFS */
-
 #ifdef __cplusplus
 namespace upb {
 class RefCounted;