More work on upbc.

4 files changed, 118 insertions, 83 deletions

diff --git a/upb_context.c b/upb_context.c
index 8025614..17df688 100644
--- a/upb_context.c
+++ b/upb_context.c
@@ -18,23 +18,24 @@ static int memrchr(char *data, char c, size_t len)
   return off;
 }
 
-bool addfd(struct upb_strtable *t, google_protobuf_FileDescriptorProto *fd);
+bool addfd(struct upb_strtable *addto, struct upb_strtable *existingdefs,
+           google_protobuf_FileDescriptorProto *fd);
 
 bool upb_context_init(struct upb_context *c)
 {
   upb_strtable_init(&c->symtab, 16, sizeof(struct upb_symtab_entry));
   upb_strtable_init(&c->psymtab, 16, sizeof(struct upb_symtab_entry));
   /* Add all the types in descriptor.proto so we can parse descriptors. */
-  if(!addfd(&c->psymtab, &google_protobuf_filedescriptor)) {
+  if(!addfd(&c->psymtab, &c->psymtab, &google_protobuf_filedescriptor)) {
     assert(false);
     return false;  /* Indicates that upb is buggy or corrupt. */
   }
-  struct upb_string name = UPB_STRLIT("google.protobuf.FileDescriptorProto");
-  c->fd_msg = upb_strtable_lookup(&c->psymtab, &name);
-  assert(c->fd_msg);
-  c->fd_size = 16;
-  c->fd_len = 0;
-  c->fd = malloc(sizeof(*c->fd));
+  struct upb_string name = UPB_STRLIT("google.protobuf.FileDescriptorSet");
+  c->fds_msg = upb_strtable_lookup(&c->psymtab, &name);
+  assert(c->fds_msg);
+  c->fds_size = 16;
+  c->fds_len = 0;
+  c->fds = malloc(sizeof(*c->fds));
   return true;
 }
 
@@ -57,8 +58,8 @@ void upb_context_free(struct upb_context *c)
 {
   free_symtab(&c->symtab);
   free_symtab(&c->psymtab);
-  for(size_t i = 0; i < c->fd_len; i++) free(c->fd[i]);
-  free(c->fd);
+  for(size_t i = 0; i < c->fds_len; i++) free(c->fds[i]);
+  free(c->fds);
 }
 
 struct upb_symtab_entry *upb_context_lookup(struct upb_context *c,
@@ -204,38 +205,29 @@ static bool insert_message(struct upb_strtable *t,
   return true;
 }
 
-bool addfd(struct upb_strtable *t, google_protobuf_FileDescriptorProto *fd)
+bool addfd(struct upb_strtable *addto, struct upb_strtable *existingdefs,
+           google_protobuf_FileDescriptorProto *fd)
 {
   struct upb_string package = {.byte_len=0};
   if(fd->set_flags.has.package) package = *fd->package;
 
-  /* We want the entire add operation to be atomic, so we initially insert into
-   * this temporary map of symbols.  Once we have verified that there are no
-   * errors (all symbols can be resolved and no illegal redefinitions occurred)
-   * only then do we insert into the context's table. */
-  struct upb_strtable tmp;
-  int symcount = (fd->set_flags.has.message_type ? fd->message_type->len : 0) +
-                 (fd->set_flags.has.enum_type ? fd->enum_type->len : 0) +
-                 (fd->set_flags.has.service ? fd->service->len : 0);
-  upb_strtable_init(&tmp, symcount, sizeof(struct upb_symtab_entry));
-
   if(fd->set_flags.has.message_type)
     for(unsigned int i = 0; i < fd->message_type->len; i++)
-      if(!insert_message(&tmp, fd->message_type->elements[i], &package))
-        goto error;
+      if(!insert_message(addto, fd->message_type->elements[i], &package))
+        return false;
 
   if(fd->set_flags.has.enum_type)
     for(unsigned int i = 0; i < fd->enum_type->len; i++)
-      if(!insert_enum(&tmp, fd->enum_type->elements[i], &package))
-        goto error;
+      if(!insert_enum(addto, fd->enum_type->elements[i], &package))
+        return false;
 
   /* TODO: handle extensions and services. */
 
   /* Attempt to resolve all references. */
   struct upb_symtab_entry *e;
-  for(e = upb_strtable_begin(&tmp); e; e = upb_strtable_next(&tmp, &e->e)) {
-    if(upb_strtable_lookup(t, &e->e.key))
-      goto error;  /* Redefinition prohibited. */
+  for(e = upb_strtable_begin(addto); e; e = upb_strtable_next(addto, &e->e)) {
+    if(upb_strtable_lookup(existingdefs, &e->e.key))
+      return false;  /* Redefinition prohibited. */
     if(e->type == UPB_SYM_MESSAGE) {
       struct upb_msg *m = e->ref.msg;
       for(unsigned int i = 0; i < m->num_fields; i++) {
@@ -243,44 +235,56 @@ bool addfd(struct upb_strtable *t, google_protobuf_FileDescriptorProto *fd)
         google_protobuf_FieldDescriptorProto *fd = m->field_descriptors[i];
         union upb_symbol_ref ref;
         if(fd->type == GOOGLE_PROTOBUF_FIELDDESCRIPTORPROTO_TYPE_MESSAGE)
-          ref = resolve2(t, &tmp, &e->e.key, fd->type_name, UPB_SYM_MESSAGE);
+          ref = resolve2(existingdefs, addto, &e->e.key, fd->type_name,
+                         UPB_SYM_MESSAGE);
         else if(fd->type == GOOGLE_PROTOBUF_FIELDDESCRIPTORPROTO_TYPE_ENUM)
-          ref = resolve2(t, &tmp, &e->e.key, fd->type_name, UPB_SYM_ENUM);
+          ref = resolve2(existingdefs, addto, &e->e.key, fd->type_name,
+                         UPB_SYM_ENUM);
         else
           continue;  /* No resolving necessary. */
-        if(!ref.msg) goto error;  /* Ref. to undefined symbol. */
+        if(!ref.msg) return false;  /* Ref. to undefined symbol. */
         upb_msg_ref(m, f, ref);
       }
     }
   }
-
-  /* All references were successfully resolved -- add to the symbol table. */
-  for(e = upb_strtable_begin(&tmp); e; e = upb_strtable_next(&tmp, &e->e))
-    upb_strtable_insert(t, &e->e);
-
-  upb_strtable_free(&tmp);
   return true;
-
-error:
-  /* TODO */
-  return false;
 }
 
 bool upb_context_addfd(struct upb_context *c,
                        google_protobuf_FileDescriptorProto *fd)
 {
-  return addfd(&c->symtab, fd);
+  struct upb_strtable tmp;
+  if(!addfd(&tmp, &c->symtab, fd)) {
+    free_symtab(&tmp);
+    return false;
+  }
+  upb_strtable_free(&tmp);
+  return true;
 }
 
-bool upb_context_parsefd(struct upb_context *c, struct upb_string *fd_str) {
-  google_protobuf_FileDescriptorProto *fd =
-      upb_alloc_and_parse(c->fd_msg, fd_str, true);
-  if(!fd) return false;
-  if(!upb_context_addfd(c, fd)) return false;
-  if(c->fd_size == c->fd_len) {
-    c->fd_size *= 2;
-    c->fd = realloc(c->fd, c->fd_size);
+bool upb_context_parsefds(struct upb_context *c, struct upb_string *fds_str) {
+  google_protobuf_FileDescriptorSet *fds =
+      upb_alloc_and_parse(c->fds_msg, fds_str, true);
+  if(!fds) return false;
+  if(fds->set_flags.has.file) {
+    struct upb_strtable tmp;
+    upb_strtable_init(&tmp, 0, sizeof(struct upb_symtab_entry));
+    for(uint32_t i = 0; i < fds->file->len; i++) {
+      if(!addfd(&tmp, &c->symtab, fds->file->elements[i])) {
+        free_symtab(&tmp);
+        return false;
+      }
+    }
+    /* Everything was successfully added, copy from the tmp symtable. */
+    struct upb_symtab_entry *e;
+    for(e = upb_strtable_begin(&tmp); e; e = upb_strtable_next(&tmp, &e->e))
+      upb_strtable_insert(&c->symtab, &e->e);
+    upb_strtable_free(&tmp);
+  }
+  if(c->fds_size == c->fds_len) {
+    c->fds_size *= 2;
+    c->fds = realloc(c->fds, c->fds_size);
   }
-  c->fd[c->fd_len++] = fd;  /* Need to keep a ref since we own it. */
+  c->fds[c->fds_len++] = fds;  /* Need to keep a ref since we own it. */
   return true;
 }
diff --git a/upb_context.h b/upb_context.h
index 8229ae4..51afed1 100644
--- a/upb_context.h
+++ b/upb_context.h
@@ -31,12 +31,12 @@ struct upb_symtab_entry {
 struct upb_context {
   struct upb_strtable symtab;   /* The context's symbol table. */
   struct upb_strtable psymtab;  /* Private symbols, for internal use. */
-  struct upb_msg *fd_msg;    /* This is in psymtab, ptr here for convenience. */
+  struct upb_msg *fds_msg;   /* This is in psymtab, ptr here for convenience. */
 
   /* A list of the FileDescriptorProtos we own (from having parsed them
    * ourselves) and must free on destruction. */
-  size_t fd_size, fd_len;
-  struct google_protobuf_FileDescriptorProto **fd;
+  size_t fds_size, fds_len;
+  struct google_protobuf_FileDescriptorSet **fds;
 };
 
 /* Initializes and frees a upb_context, respectively.  Newly initialized
@@ -93,9 +93,7 @@ INLINE struct upb_symtab_entry *upb_context_symnext(
 bool upb_context_addfd(struct upb_context *c,
                        struct google_protobuf_FileDescriptorProto *fd);
 
-/* Like the previous, but takes a serialized FileDescriptorProto and parses
- * it before adding to the context. */
-bool upb_context_parsefd(struct upb_context *c, struct upb_string *fd);
+bool upb_context_parsefds(struct upb_context *c, struct upb_string *fds);
 
 #ifdef __cplusplus
 }  /* extern "C" */
diff --git a/upb_string.h b/upb_string.h
index d91fa69..2520caa 100644
--- a/upb_string.h
+++ b/upb_string.h
@@ -59,6 +59,23 @@ INLINE void upb_strfree(struct upb_string s) {
   free(s.ptr);
 }
 
+INLINE bool upb_strreadfile(const char *filename, struct upb_string *data) {
+  FILE *f = fopen(filename, "rb");
+  if(!f) return false;
+  if(fseek(f, 0, SEEK_END) != 0) return false;
+  long size = ftell(f);
+  if(size < 0) return false;
+  if(fseek(f, 0, SEEK_SET) != 0) return false;
+  data->ptr = (char*)malloc(size);
+  data->byte_len = size;
+  if(fread(data->ptr, size, 1, f) != 1) {
+    free(data->ptr);
+    return false;
+  }
+  fclose(f);
+  return true;
+}
+
 #define UPB_STRLIT(strlit) {.ptr=strlit, .byte_len=sizeof(strlit)-1}
 #define UPB_STRARG(str) (str).byte_len, (str).ptr
 #define UPB_STRFMT "%.*s"
diff --git a/upbc.c b/upbc.c
index b49d24b..2cda5e4 100644
--- a/upbc.c
+++ b/upbc.c
@@ -10,6 +10,7 @@
 #include <inttypes.h>
 #include "descriptor.h"
 #include "upb_context.h"
+#include "upb_enum.h"
 
 /* These are in-place string transformations that do not change the length of
  * the string (and thus never need to re-allocate). */
@@ -31,12 +32,11 @@ static void to_preproc(struct upb_string str)
  * also defines constants for the enum values.
  *
  * Assumes that d has been validated. */
-static void write_header(google_protobuf_FileDescriptorProto *d, FILE *stream)
+static void write_header(struct upb_symtab_entry entries[], int num_entries,
+                         struct upb_string outfile_name, FILE *stream)
 {
   /* Header file prologue. */
-  fprintf(stream, "/* Auto-generated from " UPB_STRFMT ".  Do not edit. */\n\n",
-          UPB_STRARG(*d->name));
-  struct upb_string include_guard_name = upb_strdup(*d->name);
+  struct upb_string include_guard_name = upb_strdup(outfile_name);
   to_preproc(include_guard_name);
   fprintf(stream, "#ifndef " UPB_STRFMT "\n", UPB_STRARG(include_guard_name));
   fprintf(stream, "#define " UPB_STRFMT "\n\n", UPB_STRARG(include_guard_name));
@@ -46,29 +46,31 @@ static void write_header(google_protobuf_FileDescriptorProto *d, FILE *stream)
   fputs("#endif\n\n", stream);
 
   /* Enums. */
-  if(d->set_flags.has.enum_type) {
-    fprintf(stream, "/* Enums. */\n\n");
-    for(uint32_t i = 0; i < d->enum_type->len; i++) {  /* Foreach enum */
-      google_protobuf_EnumDescriptorProto *e = d->enum_type->elements[i];
-      struct upb_string enum_name = upb_strdup(*e->name);
-      to_cident(enum_name);
-      fprintf(stream, "typedef enum " UPB_STRFMT " {\n", UPB_STRARG(enum_name));
-      if(e->set_flags.has.value) {
-        for(uint32_t j = 0; j < e->value->len; j++) {  /* Foreach enum value. */
-          google_protobuf_EnumValueDescriptorProto *v = e->value->elements[i];
-          struct upb_string value_name = upb_strdup(*v->name);
-          to_preproc(value_name);
-          /* "  GOOGLE_PROTOBUF_FIELDDESCRIPTORPROTO_TYPE_UINT32 = 13," */
-          fprintf(stream, "  " UPB_STRFMT " = %" PRIu32,
-                  UPB_STRARG(value_name), v->number);
-          if(j != e->value->len-1) fputc(',', stream);
-          fputc('\n', stream);
-          upb_strfree(value_name);
-        }
+  fprintf(stream, "/* Enums. */\n\n");
+  for(int i = 0; i < num_entries; i++) {  /* Foreach enum */
+    if(entries[i].type != UPB_SYM_ENUM) continue;
+    struct upb_symtab_entry *entry = &entries[i];
+    struct upb_enum *e = entry->ref._enum;
+    google_protobuf_EnumDescriptorProto *ed = e->descriptor;
+    /* We use entry->e.key (the fully qualified name) instead of ed->name. */
+    struct upb_string enum_name = upb_strdup(entry->e.key);
+    to_cident(enum_name);
+    fprintf(stream, "typedef enum " UPB_STRFMT " {\n", UPB_STRARG(enum_name));
+    if(ed->set_flags.has.value) {
+      for(uint32_t j = 0; j < ed->value->len; j++) {  /* Foreach enum value. */
+        google_protobuf_EnumValueDescriptorProto *v = ed->value->elements[i];
+        struct upb_string value_name = upb_strdup(*v->name);
+        to_preproc(value_name);
+        /* "  GOOGLE_PROTOBUF_FIELDDESCRIPTORPROTO_TYPE_UINT32 = 13," */
+        fprintf(stream, "  " UPB_STRFMT " = %" PRIu32,
+                UPB_STRARG(value_name), v->number);
+        if(j != ed->value->len-1) fputc(',', stream);
+        fputc('\n', stream);
+        upb_strfree(value_name);
       }
-      fprintf(stream, "} " UPB_STRFMT ";\n\n", UPB_STRARG(enum_name));
-      upb_strfree(enum_name);
     }
+    fprintf(stream, "} " UPB_STRFMT ";\n\n", UPB_STRARG(enum_name));
+    upb_strfree(enum_name);
   }
 
   /* Epilogue. */
@@ -81,6 +83,20 @@ static void write_header(google_protobuf_FileDescriptorProto *d, FILE *stream)
 
 int main()
 {
-  write_header(&google_protobuf_filedescriptor, stdout);
+  struct upb_context c;
+  upb_context_init(&c);
+  struct upb_string fds;
+  assert(upb_strreadfile("/tmp/descriptor.proto.bin", &fds));
+  assert(upb_context_parsefds(&c, &fds));
+  struct upb_strtable *t = &c.symtab;
+  int symcount = t->t.count;
+  struct upb_symtab_entry entries[symcount];
+  struct upb_symtab_entry *e = upb_strtable_begin(t);
+  int i = 0;
+  for(; e && i < symcount; e = upb_strtable_next(t, &e->e), i++)
+    entries[i] = *e;
+  assert(e == NULL && i == symcount);
+  struct upb_string name = UPB_STRLIT("descriptor.proto");
+  write_header(entries, symcount, name, stdout);
 }