Reference-count upb_msgdef and upb_enumdef.

The context owns a reference on each def, defs own references
on defs they reference, and msgs own refs on their def.

src/upb_context.c

@@ -56,11 +56,10 @@ static void free_symtab(struct upb_strtable *t)
   struct upb_symtab_entry *e = upb_strtable_begin(t);
   for(; e; e = upb_strtable_next(t, &e->e)) {
     switch(e->type) {
-      case UPB_SYM_MESSAGE: upb_msgdef_free(e->ref.msg); break;
-      case UPB_SYM_ENUM: upb_enumdef_free(e->ref._enum); break;
+      case UPB_SYM_MESSAGE: upb_msgdef_unref(e->ref.msg); break;
+      case UPB_SYM_ENUM: upb_enumdef_unref(e->ref._enum); break;
       default: break;  /* TODO */
     }
-    free(e->ref.msg);  /* The pointer is the same for all. */
     free(e->e.key.ptr);
   }
   upb_strtable_free(t);

src/upb_def.c

@@ -47,6 +47,7 @@ void upb_msgdef_init(struct upb_msgdef *m, google_protobuf_DescriptorProto *d,
   upb_strtable_init(&m->fields_by_name, num_fields,
                     sizeof(struct upb_fieldsbyname_entry));
 
+  upb_atomic_refcount_init(&m->refcount, 1);
   m->fqname = upb_strdup(fqname);
   m->context = c;
   m->num_fields = num_fields;
@@ -101,15 +102,18 @@ void upb_msgdef_init(struct upb_msgdef *m, google_protobuf_DescriptorProto *d,
   free(fds);
 }
 
-void upb_msgdef_free(struct upb_msgdef *m)
+void _upb_msgdef_free(struct upb_msgdef *m)
 {
   upb_inttable_free(&m->fields_by_num);
   upb_strtable_free(&m->fields_by_name);
   upb_string_unref(m->fqname);
   for (unsigned int i = 0; i < m->num_fields; i++) {
-    upb_string_unref(m->fields[i].name);
+    struct upb_fielddef *f = &m->fields[i];
+    upb_string_unref(f->name);
+    upb_def_unref(f->ref, f->type);
   }
   free(m->fields);
+  free(m);
 }
 
 void upb_msgdef_setref(struct upb_msgdef *m, struct upb_fielddef *f,
@@ -122,15 +126,15 @@ void upb_msgdef_setref(struct upb_msgdef *m, struct upb_fielddef *f,
   f->ref = ref;
   int_e->f.ref = ref;
   str_e->f.ref = ref;
+  upb_def_ref(ref, f->type);
 }
 
-
 void upb_enumdef_init(struct upb_enumdef *e,
                    struct google_protobuf_EnumDescriptorProto *ed,
                    struct upb_context *c) {
   int num_values = ed->set_flags.has.value ? ed->value->len : 0;
   e->context = c;
-  upb_atomic_refcount_init(&e->refcount, 0);
+  upb_atomic_refcount_init(&e->refcount, 1);
   upb_strtable_init(&e->nametoint, num_values,
                     sizeof(struct upb_enumdef_ntoi_entry));
   upb_inttable_init(&e->inttoname, num_values,
@@ -147,7 +151,8 @@ void upb_enumdef_init(struct upb_enumdef *e,
   }
 }
 
-void upb_enumdef_free(struct upb_enumdef *e) {
+void _upb_enumdef_free(struct upb_enumdef *e) {
   upb_strtable_free(&e->nametoint);
   upb_inttable_free(&e->inttoname);
+  free(e);
 }

src/upb_def.h

@@ -9,7 +9,11 @@
  * - upb_enumdef: describes an enum.
  * (TODO: descriptions of extensions and services).
  *
- * Defs are immutable and reference-counted.  
+ * Defs are immutable and reference-counted.  Contexts reference any defs
+ * that are the currently active def for that context, but they can be
+ * unref'd if the message is changed by loading extensions.  In the case
+ * that a def is no longer active in a context, it will still be ref'd by
+ * messages (if any) that were constructed with that def.
  *
  * This file contains routines for creating and manipulating the definitions
  * themselves.  To create and manipulate actual messages, see upb_msg.h.
@@ -32,8 +36,10 @@ struct upb_context;
 struct google_protobuf_EnumDescriptorProto;
 struct google_protobuf_DescriptorProto;
 struct google_protobuf_FieldDescriptorProto;
+
 /* Structure that describes a single .proto message type. */
 struct upb_msgdef {
+  upb_atomic_refcount_t refcount;
   struct upb_context *context;
   struct upb_msg *default_msg;   /* Message with all default values set. */
   struct upb_string *fqname;     /* Fully qualified. */
@@ -149,20 +155,12 @@ struct upb_enumdef_iton_entry {
   struct upb_string *string;
 };
 
-/* Initializes and frees an enum, respectively.  Caller retains ownership of
- * ed, but it must outlive e. */
-void upb_enumdef_init(struct upb_enumdef *e,
-                      struct google_protobuf_EnumDescriptorProto *ed,
-                      struct upb_context *c);
-void upb_enumdef_free(struct upb_enumdef *e);
-
-
 /* Internal functions. ********************************************************/
 
 /* Initializes/frees a upb_msgdef.  Usually this will be called by upb_context,
  * and clients will not have to construct one directly.
  *
- * Caller retains ownership of d.  Note that init does not resolve
+ * Caller retains ownership of d and fqname.  Note that init does not resolve
  * upb_fielddef.ref the caller should do that post-initialization by
  * calling upb_msg_ref() below.
  *
@@ -175,7 +173,13 @@ void upb_msgdef_init(struct upb_msgdef *m,
                      struct google_protobuf_DescriptorProto *d,
                      struct upb_string *fqname, bool sort,
                      struct upb_context *c, struct upb_status *status);
-void upb_msgdef_free(struct upb_msgdef *m);
+void _upb_msgdef_free(struct upb_msgdef *m);
+INLINE void upb_msgdef_ref(struct upb_msgdef *m) {
+  upb_atomic_ref(&m->refcount);
+}
+INLINE void upb_msgdef_unref(struct upb_msgdef *m) {
+  if(upb_atomic_unref(&m->refcount)) _upb_msgdef_free(m);
+}
 
 /* Sort the given field descriptors in-place, according to what we think is an
  * optimal ordering of fields.  This can change from upb release to upb
@@ -190,6 +194,47 @@ void upb_msgdef_sortfds(struct google_protobuf_FieldDescriptorProto **fds,
 void upb_msgdef_setref(struct upb_msgdef *m, struct upb_fielddef *f,
                        union upb_symbol_ref ref);
 
+/* Initializes and frees an enum, respectively.  Caller retains ownership of
+ * ed.  The enumdef is initialized with one ref. */
+void upb_enumdef_init(struct upb_enumdef *e,
+                      struct google_protobuf_EnumDescriptorProto *ed,
+                      struct upb_context *c);
+void _upb_enumdef_free(struct upb_enumdef *e);
+INLINE void upb_enumdef_ref(struct upb_enumdef *e) {
+  upb_atomic_ref(&e->refcount);
+}
+INLINE void upb_enumdef_unref(struct upb_enumdef *e) {
+  if(upb_atomic_unref(&e->refcount)) _upb_enumdef_free(e);
+}
+
+INLINE void upb_def_ref(union upb_symbol_ref ref, upb_field_type_t type) {
+  switch(type) {
+    case UPB_TYPENUM(MESSAGE):
+    case UPB_TYPENUM(GROUP):
+      upb_msgdef_ref(ref.msg);
+      break;
+    case UPB_TYPENUM(ENUM):
+      upb_enumdef_ref(ref._enum);
+      break;
+    default:
+      assert(false);
+  }
+}
+
+INLINE void upb_def_unref(union upb_symbol_ref ref, upb_field_type_t type) {
+  switch(type) {
+    case UPB_TYPENUM(MESSAGE):
+    case UPB_TYPENUM(GROUP):
+      upb_msgdef_unref(ref.msg);
+      break;
+    case UPB_TYPENUM(ENUM):
+      upb_enumdef_unref(ref._enum);
+      break;
+    default:
+      assert(false);
+  }
+}
+
 #ifdef __cplusplus
 }  /* extern "C" */
 #endif

src/upb_mm.c

@@ -24,6 +24,7 @@ void upb_msg_destroy(struct upb_msg *msg) {
     if(!upb_msg_isset(msg, f) || !upb_field_ismm(f)) continue;
     upb_mm_destroy(upb_msg_getptr(msg, f), upb_field_ptrtype(f));
   }
+  upb_msgdef_unref(msg->def);
   free(msg);
 }
 

src/upb_msg.h

@@ -59,6 +59,7 @@ INLINE struct upb_msg *upb_msg_new(struct upb_msgdef *md) {
   memset(msg, 0, size);
   upb_mmhead_init(&msg->mmhead);
   msg->def = md;
+  upb_msgdef_ref(md);
   return msg;
 }