Other than a couple memory leaks, "make descriptorgen" works again!

4 files changed, 47 insertions, 31 deletions

diff --git a/src/upb_data.c b/src/upb_data.c
index 084a03c..8de0c62 100644
--- a/src/upb_data.c
+++ b/src/upb_data.c
@@ -41,7 +41,7 @@ static void data_unref(void *d, struct upb_fielddef *f) {
 }
 
 INLINE void data_init(upb_data *d, int flags) {
-  d->v = flags;
+  d->v = REFCOUNT_ONE | flags;
 }
 
 static void check_not_frozen(upb_data *d) {
@@ -100,7 +100,10 @@ void upb_string_resize(upb_string *s, upb_strlen_t byte_len) {
 
 upb_string *upb_string_getref(upb_string *s, int ref_flags) {
   if(_upb_data_incref(&s->common.base, ref_flags)) return s;
-  return upb_strdup(s);
+  upb_string *copy = upb_strdup(s);
+  if(ref_flags == UPB_REF_FROZEN)
+    upb_data_setflag(&copy->common.base, UPB_DATA_FROZEN);
+  return copy;
 }
 
 upb_string *upb_strreadfile(const char *filename) {
@@ -200,13 +203,18 @@ static void array_set_size(upb_array *a, upb_arraylen_t newsize) {
   }
 }
 
-void upb_array_resize(upb_array *a, upb_strlen_t len) {
+void upb_array_resize(upb_array *a, struct upb_fielddef *f, upb_strlen_t len) {
   check_not_frozen(&a->common.base);
-  if(array_get_size(a) < len) {
+  size_t type_size = upb_type_info[f->type].size;
+  upb_arraylen_t old_size = array_get_size(a);
+  if(old_size < len) {
     // Need to resize.
     size_t new_size = round_up_to_pow2(len);
-    a->common.elements._void = realloc(a->common.elements._void, new_size);
+    a->common.elements._void = realloc(a->common.elements._void, new_size * type_size);
     array_set_size(a, new_size);
+    memset(UPB_INDEX(a->common.elements._void, old_size, type_size),
+           0,
+           (new_size - old_size) * type_size);
   }
   a->common.len = len;
 }
@@ -269,9 +277,11 @@ static union upb_value_ptr get_value_ptr(upb_msg *msg, struct upb_fielddef *f)
       }
       upb_array_truncate(*p.arr);
       upb_msg_sethas(msg, f);
+    } else {
+      assert(*p.arr);
     }
     upb_arraylen_t oldlen = upb_array_len(*p.arr);
-    upb_array_resize(*p.arr, oldlen + 1);
+    upb_array_resize(*p.arr, f, oldlen + 1);
     p = _upb_array_getptr(*p.arr, f, oldlen);
   }
   return p;
diff --git a/src/upb_data.h b/src/upb_data.h
index 872f237..2a7940e 100644
--- a/src/upb_data.h
+++ b/src/upb_data.h
@@ -130,8 +130,9 @@ typedef enum {
 // Attempts to increment the reference on d with the given type of ref.  If
 // this is not possible, returns false.
 INLINE bool _upb_data_incref(upb_data *d, upb_reftype reftype) {
-  if((reftype == UPB_REF_FROZEN && !upb_data_hasflag(d, UPB_DATA_FROZEN)) ||
-     (reftype == UPB_REF_MUTABLE && upb_data_hasflag(d, UPB_DATA_FROZEN)) ||
+  bool frozen = upb_data_hasflag(d, UPB_DATA_FROZEN);
+  if((reftype == UPB_REF_FROZEN && !frozen) ||
+     (reftype == UPB_REF_MUTABLE && frozen) ||
      (upb_data_hasflag(d, UPB_DATA_HEAPALLOCATED) &&
       !upb_data_hasflag(d, UPB_DATA_REFCOUNTED))) {
     return false;
@@ -306,7 +307,7 @@ upb_string *upb_strreadfile(const char *filename);
 // must not dynamically allocate this type.
 typedef upb_string upb_static_string;
 #define UPB_STRLIT_LEN(str, len) {0 | UPB_DATA_FROZEN, len, str}
-#define UPB_STRLIT(str) {{0 | UPB_DATA_FROZEN, sizeof(str), str}}
+#define UPB_STRLIT(str) {{0 | UPB_DATA_FROZEN, sizeof(str)-1, str}}
 
 // Allows using upb_strings in printf, ie:
 //   upb_string str = UPB_STRLIT("Hello, World!\n");
@@ -359,6 +360,11 @@ typedef struct type ## _array { \
 // empty.  Caller owns one ref on it.
 upb_array *upb_array_new(void);
 
+// Returns the current number of elements in the array.
+INLINE size_t upb_array_len(upb_array *a) {
+  return a->common.len;
+}
+
 // INTERNAL-ONLY:
 // Frees the given message and releases references on members.
 void _upb_array_free(upb_array *a, struct upb_fielddef *f);
@@ -366,7 +372,8 @@ void _upb_array_free(upb_array *a, struct upb_fielddef *f);
 // INTERNAL-ONLY:
 // Returns a pointer to the given elem.
 INLINE union upb_value_ptr _upb_array_getptr(upb_array *a,
-                                             struct upb_fielddef *f, int elem) {
+                                             struct upb_fielddef *f,
+                                             upb_arraylen_t elem) {
   assert(elem < upb_array_len(a));
   size_t type_size = upb_type_info[f->type].size;
   union upb_value_ptr p = {._void = &a->common.elements.uint8[elem * type_size]};
@@ -405,11 +412,6 @@ INLINE void upb_array_append_default(upb_array *a, struct upb_fielddef *f,
                              union upb_value val);
 #endif
 
-// Returns the current number of elements in the array.
-INLINE size_t upb_array_len(upb_array *a) {
-  return a->common.len;
-}
-
 INLINE void upb_array_truncate(upb_array *a) {
   a->common.len = 0;
 }
@@ -447,7 +449,7 @@ INLINE void upb_msg_unref(upb_msg *msg, struct upb_msgdef *md) {
 // Tests whether the given field is explicitly set, or whether it will return
 // a default.
 INLINE bool upb_msg_has(upb_msg *msg, struct upb_fielddef *f) {
-  return msg->data[f->field_index/8] % (1 << (f->field_index % 8));
+  return (msg->data[f->field_index/8] & (1 << (f->field_index % 8))) != 0;
 }
 
 // Returns the current value if set, or the default value if not set, of the
diff --git a/src/upb_def.c b/src/upb_def.c
index be5c181..f824910 100644
--- a/src/upb_def.c
+++ b/src/upb_def.c
@@ -287,7 +287,7 @@ static struct upb_msgdef *msgdef_new(struct upb_fielddef **fields,
   m->set_flags_bytes = div_round_up(m->num_fields, 8);
   // These are incremented in the loop.
   m->num_required_fields = 0;
-  m->size = m->set_flags_bytes;
+  m->size = m->set_flags_bytes + 4;  // 4 for the refcount.
   m->fields = malloc(sizeof(struct upb_fielddef) * num_fields);
 
   size_t max_align = 0;
@@ -300,8 +300,9 @@ static struct upb_msgdef *msgdef_new(struct upb_fielddef **fields,
     // multiple of that type's alignment.  Also, the size of the structure as
     // a whole must be a multiple of the greatest alignment of any member. */
     f->field_index = i;
-    f->byte_offset = ALIGN_UP(m->size, type_info->align);
-    m->size = f->byte_offset + type_info->size;
+    size_t offset = ALIGN_UP(m->size, type_info->align);
+    f->byte_offset = offset - 4;  // Offsets are relative to the refcount.
+    m->size = offset + type_info->size;
     max_align = UPB_MAX(max_align, type_info->align);
     if(f->label == UPB_LABEL(REQUIRED)) {
       // We currently rely on the fact that required fields are always sorted
@@ -516,7 +517,7 @@ static void insert_message(struct upb_strtable *t,
 
   int num_fields = d->set_flags.has.field ? d->field->len : 0;
   struct symtab_ent e;
-  e.e.key = fqname;  // Donating our ref to the table.
+  e.e.key = fqname;
   struct upb_fielddef **fielddefs = malloc(sizeof(*fielddefs) * num_fields);
   for (int i = 0; i < num_fields; i++) {
     google_protobuf_FieldDescriptorProto *fd = d->field->elements[i];
@@ -530,7 +531,7 @@ static void insert_message(struct upb_strtable *t,
   }
   free(fielddefs);
 
-  if(!upb_ok(status)) return;
+  if(!upb_ok(status)) goto error;
 
   upb_strtable_insert(t, &e.e);
 
@@ -542,6 +543,9 @@ static void insert_message(struct upb_strtable *t,
   if(d->set_flags.has.enum_type)
     for(unsigned int i = 0; i < d->enum_type->len; i++)
       insert_enum(t, d->enum_type->elements[i], fqname, status);
+
+error:
+  upb_string_unref(fqname);
 }
 
 static bool find_cycles(struct upb_msgdef *m, int search_depth,
diff --git a/tools/upbc.c b/tools/upbc.c
index 5d90fcc..f2665f2 100644
--- a/tools/upbc.c
+++ b/tools/upbc.c
@@ -136,7 +136,7 @@ static void write_h(struct upb_def *defs[], int num_defs, char *outfile_name,
         "Do not edit. */\n\n", stream),
   fprintf(stream, "#ifndef " UPB_STRFMT "\n", UPB_STRARG(include_guard_name));
   fprintf(stream, "#define " UPB_STRFMT "\n\n", UPB_STRARG(include_guard_name));
-  fputs("#include <upb_struct.h>\n\n", stream);
+  fputs("#include <upb_data.h>\n\n", stream);
   fputs("#ifdef __cplusplus\n", stream);
   fputs("extern \"C\" {\n", stream);
   fputs("#endif\n\n", stream);
@@ -171,8 +171,7 @@ static void write_h(struct upb_def *defs[], int num_defs, char *outfile_name,
     upb_string *msg_name = upb_strdup(UPB_UPCAST(m)->fqname);
     to_cident(msg_name);
     fprintf(stream, "struct " UPB_STRFMT " {\n", UPB_STRARG(msg_name));
-    fputs("  struct upb_mmhead mmhead;\n", stream);
-    fputs("  struct upb_msgdef *def;\n", stream);
+    fputs("  upb_data base;\n", stream);
     fputs("  union {\n", stream);
     fprintf(stream, "    uint8_t bytes[%" PRIu32 "];\n", m->set_flags_bytes);
     fputs("    struct {\n", stream);
@@ -214,8 +213,8 @@ static void write_h(struct upb_def *defs[], int num_defs, char *outfile_name,
       } else {
         static char* c_types[] = {
           "", "double", "float", "int64_t", "uint64_t", "int32_t", "uint64_t",
-          "uint32_t", "bool", "struct upb_string*", "", "",
-          "struct upb_string*", "uint32_t", "int32_t", "int32_t", "int64_t",
+          "uint32_t", "bool", "upb_string*", "", "",
+          "upb_string*", "uint32_t", "int32_t", "int32_t", "int64_t",
           "int32_t", "int64_t"
         };
         fprintf(stream, "  %s " UPB_STRFMT ";\n",
@@ -450,10 +449,10 @@ static void write_message_c(upb_msg *msg, struct upb_msgdef *md,
   }
   fputs("\";\n\n", stream);
 
-  fputs("static struct upb_string strings[] = {\n", stream);
+  fputs("static upb_norefcount_string strings[] = {\n", stream);
   for(int i = 0; i < size; i++) {
     struct strtable_entry *e = str_entries[i];
-    fprintf(stream, "  {.ptr = &strdata[%d], .byte_len=%d},\n", e->offset, upb_strlen(e->e.key));
+    fprintf(stream, "  UPB_STRLIT_LEN(&strdata[%d], %d),\n", e->offset, upb_strlen(e->e.key));
   }
   fputs("};\n\n", stream);
   free(str_entries);
@@ -509,8 +508,9 @@ static void write_message_c(upb_msg *msg, struct upb_msgdef *md,
       if(upb_issubmsg(e->field)) {
         struct upb_msgdef *m = upb_downcast_msgdef(e->field->def);
         void *msgdata = val.msg;
+        fputs("  {.base = {UPB_DATA_FROZEN},\n", stream);
         /* Print set flags. */
-        fputs("  {.set_flags = {.has = {\n", stream);
+        fputs("   .set_flags = {.has = {\n", stream);
         for(upb_field_count_t j = 0; j < m->num_fields; j++) {
           struct upb_fielddef *f = &m->fields[j];
           fprintf(stream, "    ." UPB_STRFMT " = ", UPB_STRARG(f->name));
@@ -535,7 +535,7 @@ static void write_message_c(upb_msg *msg, struct upb_msgdef *md,
             } else {
               struct strtable_entry *str_e = upb_strtable_lookup(&strings, val.str);
               assert(str_e);
-              fprintf(stream, "&strings[%d],   /* \"" UPB_STRFMT "\" */",
+              fprintf(stream, "(upb_string*)&strings[%d],   /* \"" UPB_STRFMT "\" */",
                       str_e->num, UPB_STRARG(val.str));
             }
           } else if(upb_isarray(f)) {
@@ -717,7 +717,6 @@ int main(int argc, char *argv[])
 
   int symcount;
   struct upb_def **defs = upb_symtab_getdefs(s, &symcount, UPB_DEF_ANY);
-  upb_symtab_unref(s);
   write_h(defs, symcount, h_filename, cident, h_file);
   write_const_h(defs, symcount, h_filename, h_const_file);
   for (int i = 0; i < symcount; i++) upb_def_unref(defs[i]);
@@ -730,6 +729,7 @@ int main(int argc, char *argv[])
   }
   upb_msg_unref(fds_msg, s->fds_msgdef);
   upb_string_unref(descriptor);
+  upb_symtab_unref(s);
   fclose(h_file);
   fclose(h_const_file);