JIT passes all tests!

5 files changed, 77 insertions, 46 deletions

diff --git a/src/upb_decoder.c b/src/upb_decoder.c
index 4b71ccd..3107f77 100644
--- a/src/upb_decoder.c
+++ b/src/upb_decoder.c
@@ -67,7 +67,7 @@ static bool upb_pullbuf(upb_decoder *d) {
   d->buf = upb_string_getrobuf(d->bufstr);
   d->ptr = upb_string_getrobuf(d->bufstr);
   d->end = d->buf + upb_string_len(d->bufstr);
-  d->jit_end = d->end; //d->end - 12;
+  d->jit_end = d->end - 20;
   upb_string_substr(d->tmp, d->bufstr, 0, 0);
   upb_dstate_setmsgend(d);
   return true;
@@ -247,10 +247,13 @@ void upb_decoder_decode(upb_decoder *d, upb_status *status) {
     // before falling through to the slow(er) path.
 #ifdef UPB_USE_JIT_X64
     void (*upb_jit_decode)(upb_decoder *d) = (void*)d->jit_code;
-    if (d->dispatcher.handlers->should_jit && d->buf) {
-      //fprintf(stderr, "Entering JIT, ptr: %p\n", d->ptr);
+    if (d->jit_code && d->dispatcher.top == d->dispatcher.stack && d->ptr < d->jit_end) {
+      const char *before = d->ptr;
+      //fprintf(stderr, "Entering JIT, JIT bytes left: %zd\n", d->jit_end - d->ptr);
       upb_jit_decode(d);
-      //fprintf(stderr, "Exiting JIT, ptr: %p\n", d->ptr);
+      //fprintf(stderr, "Exiting JIT, parsed %zd bytes\n", d->ptr - before);
+      //fprintf(stderr, "ptr: %p, effective_end: %p, jit_end: %p, effective_end-ptr=%d\n",
+      //        d->ptr, d->effective_end, d->jit_end, d->effective_end - d->ptr);
     }
 #endif
 
@@ -355,7 +358,8 @@ err:
 void upb_decoder_init(upb_decoder *d, upb_handlers *handlers) {
   upb_dispatcher_init(&d->dispatcher, handlers);
 #ifdef UPB_USE_JIT_X64
-  upb_decoder_makejit(d);
+  d->jit_code = NULL;
+  if (d->dispatcher.handlers->should_jit) upb_decoder_makejit(d);
 #endif
   d->bufstr = NULL;
   d->buf = NULL;
@@ -378,6 +382,6 @@ void upb_decoder_uninit(upb_decoder *d) {
   upb_string_unref(d->bufstr);
   upb_string_unref(d->tmp);
 #ifdef UPB_USE_JIT_X64
-  upb_decoder_freejit(d);
+  if (d->dispatcher.handlers->should_jit) upb_decoder_freejit(d);
 #endif
 }
diff --git a/src/upb_decoder_x86.dasc b/src/upb_decoder_x86.dasc
index 71df08f..ecd9899 100644
--- a/src/upb_decoder_x86.dasc
+++ b/src/upb_decoder_x86.dasc
@@ -88,7 +88,7 @@ void __attribute__((noinline)) __jit_debug_register_code() { __asm__ __volatile_
 |// Checks PTR for end-of-buffer.
 |.macro check_eob, m
 |  cmp   PTR, DECODER->effective_end
-|| if (m->is_group) {
+|| if (m->endgroup_f) {
      |  jae  ->exit_jit
 || } else {
      |  jae  =>m->jit_endofbuf_pclabel
@@ -136,13 +136,13 @@ void __attribute__((noinline)) __jit_debug_register_code() { __asm__ __volatile_
 |  and  edx, 0x7
 |  cmp  ecx, m->max_field_number  // Bounds-check the field.
 |  ja   ->exit_jit                // In the future; could be unknown label
-|  mov  rcx, qword [rcx*8 + m->tablearray]  // TODO: support hybrid array/hash tables.
-|  jmp  rcx  // Dispatch: unpredictable jump.
+|  mov  rax, qword [rcx*8 + m->tablearray]  // TODO: support hybrid array/hash tables.
+|  jmp  rax  // Dispatch: unpredictable jump.
 |.endmacro
 |
 |.macro setmsgend, m
 |    mov    rsi, DECODER->jit_end
-|| if (m->is_group) {
+|| if (m->endgroup_f) {
 |    mov64  rax, 0xffffffffffffffff
 |    mov    qword DECODER->submsg_end, rax
 |    mov    DECODER->effective_end, rsi
@@ -153,9 +153,9 @@ void __attribute__((noinline)) __jit_debug_register_code() { __asm__ __volatile_
 |    add    rax, qword DECODER->buf
 |    mov    DECODER->submsg_end, rax  // submsg_end = d->buf + f->end_offset
 |    cmp    rax, rsi
-|    jb     >1
+|    jb     >8
 |    mov    rax, rsi                  // effective_end = min(d->submsg_end, d->jit_end)
-|1:
+|8:
 |    mov    DECODER->effective_end, rax
 || }
 |.endmacro
@@ -228,8 +228,8 @@ static void upb_decoder_jit_field(upb_decoder *d, uint32_t tag, uint32_t next_ta
   // We check the wire type (which must be loaded in edx) because the
   // table is keyed on field number, not type.
   |=>f->jit_pclabel:
-  |  cmp  edx, upb_types[f->type].native_wire_type
-  |  jne  ->exit_jit     // In the future: could be an unknown field.
+  |  cmp  edx, (tag & 0x7)
+  |  jne  ->exit_jit     // In the future: could be an unknown field or packed.
   |=>f->jit_pclabel_notypecheck:
   |1:  // Label for repeating this field.
 
@@ -305,10 +305,12 @@ static void upb_decoder_jit_field(upb_decoder *d, uint32_t tag, uint32_t next_ta
       |  jmp  =>m->jit_endofmsg_pclabel
       return;
 
+    // Will dispatch callbacks and call submessage in a second.
     case UPB_TYPE(MESSAGE):
       |  decode_varint  tag_size
+      break;
     case UPB_TYPE(GROUP):
-      // Will dispatch callbacks and call submessage in a second.
+      |  add  PTR, tag_size
       break;
 
     default: abort();
@@ -340,10 +342,15 @@ static void upb_decoder_jit_field(upb_decoder *d, uint32_t tag, uint32_t next_ta
     |   jae   ->exit_jit  // Frame stack overflow.
     |   mov   qword FRAME:rax->f, f
     |   mov   qword FRAME:rax->closure, rdx
-    |   mov   rsi, PTR
-    |   sub   rsi, DECODER->buf
-    |   add   r12d, esi
-    |   mov   dword FRAME:rax->end_offset, r12d   // = (d->ptr - d->buf) + delim_len
+    if (f->type == UPB_TYPE(MESSAGE)) {
+      |   mov   rsi, PTR
+      |   sub   rsi, DECODER->buf
+      |   add   r12d, esi
+      |   mov   dword FRAME:rax->end_offset, r12d   // = (d->ptr - d->buf) + delim_len
+    } else {
+      assert(f->type == UPB_TYPE(GROUP));
+      |   mov   dword FRAME:rax->end_offset, -1U
+    }
     |   mov   CLOSURE, rdx
     |   mov   DECODER->dispatcher.top, rax
     |   mov   FRAME, rax
@@ -440,11 +447,10 @@ static void upb_decoder_jit_msg(upb_decoder *d, upb_handlers_msgent *m) {
 
   free(keys);
 
-  if (m->is_group) {
-    // Create a fake fieldent for handling "end group."
-    upb_handlers_fieldent f = {0, UPB_TYPE_ENDGROUP, 0, UPB_NO_VALUE, {NULL}, NULL, 0, 0, 0, false};
-    upb_decoder_jit_field(d, last_tag, m->groupnum, m, last_f, &f);
-    upb_decoder_jit_field(d, m->groupnum, 0, m, &f, NULL);
+  if (m->endgroup_f) {
+    uint32_t tag = m->endgroup_f->number << 3 | UPB_WIRE_TYPE_END_GROUP;
+    upb_decoder_jit_field(d, last_tag, tag, m, last_f, m->endgroup_f);
+    upb_decoder_jit_field(d, tag, 0, m, m->endgroup_f, NULL);
   } else {
     upb_decoder_jit_field(d, last_tag, 0, m, last_f, NULL);
   }
@@ -452,7 +458,7 @@ static void upb_decoder_jit_msg(upb_decoder *d, upb_handlers_msgent *m) {
   // --------- New code section (does not fall through) ------------------------
 
   // End-of-buf / end-of-message.
-  if (!m->is_group) {
+  if (!m->endgroup_f) {
     // This case doesn't exist for groups, because there eob really means
     // eob, so that case just exits the jit directly.
     |=>m->jit_endofbuf_pclabel:
@@ -479,6 +485,9 @@ static void upb_decoder_jit_msg(upb_decoder *d, upb_handlers_msgent *m) {
 
 }
 
+static const char *dbgfmt =
+    "JIT encountered unknown field!  wt=%d, fn=%d\n";
+
 static void upb_decoder_jit(upb_decoder *d) {
   |  push  rbp
   |  mov   rbp, rsp
@@ -512,7 +521,10 @@ static void upb_decoder_jit(upb_decoder *d) {
   |  leave
   |  ret
   |=>0:
-  |  callp &abort
+  |  mov rdi, stderr
+  |  mov rsi, dbgfmt
+  |  callp  fprintf
+  |  callp  abort
 }
 
 void upb_decoder_jit_assignfieldlabs(upb_handlers_fieldent *f,
@@ -522,7 +534,8 @@ void upb_decoder_jit_assignfieldlabs(upb_handlers_fieldent *f,
   f->jit_submsg_done_pclabel = (*pclabel_count)++;
 }
 
-void upb_decoder_jit_assignmsglabs(upb_handlers_msgent *m,
+void upb_decoder_jit_assignmsglabs(upb_handlers *h,
+                                   upb_handlers_msgent *m,
                                    uint32_t *pclabel_count) {
   m->jit_startmsg_pclabel = (*pclabel_count)++;
   m->jit_endofbuf_pclabel = (*pclabel_count)++;
@@ -537,9 +550,17 @@ void upb_decoder_jit_assignmsglabs(upb_handlers_msgent *m,
     m->max_field_number = UPB_MAX(m->max_field_number, key);
     upb_handlers_fieldent *f = upb_inttable_iter_value(i);
     upb_decoder_jit_assignfieldlabs(f, pclabel_count);
+    if (f->type == UPB_TYPE(GROUP)) {
+      upb_handlers_msgent *sub_m = upb_handlers_getmsgent(h, f);
+      sub_m->endgroup_f = malloc(sizeof(*sub_m->endgroup_f));
+      memcpy(sub_m->endgroup_f, f, sizeof(*f));
+      sub_m->endgroup_f->type = UPB_TYPE_ENDGROUP;
+      upb_decoder_jit_assignfieldlabs(sub_m->endgroup_f, pclabel_count);
+    }
   }
   // XXX: Won't work for large field numbers; will need to use a upb_table.
-  m->tablearray = malloc((m->max_field_number + 1) * sizeof(void*));
+  // +2 to cover group case, in case group number is larger than all tags.
+  m->tablearray = malloc((m->max_field_number + 2) * sizeof(void*));
 }
 
 // Second pass: for messages that have only one parent, link them to the field
@@ -551,10 +572,6 @@ void upb_decoder_jit_assignmsglabs2(upb_handlers *h, upb_handlers_msgent *m) {
     upb_handlers_fieldent *f = upb_inttable_iter_value(i);
     if (upb_issubmsgtype(f->type)) {
       upb_handlers_msgent *sub_m = upb_handlers_getmsgent(h, f);
-      if (f->type == UPB_TYPE(GROUP)) {
-        sub_m->is_group = true;
-        sub_m->groupnum = upb_inttable_iter_key(i);
-      }
       if (sub_m->jit_parent_field_done_pclabel == UPB_NONE) {
         sub_m->jit_parent_field_done_pclabel = f->jit_submsg_done_pclabel;
       } else {
@@ -569,7 +586,7 @@ void upb_decoder_makejit(upb_decoder *d) {
   uint32_t pclabel_count = 1;
   upb_handlers *h = d->dispatcher.handlers;
   for (int i = 0; i < h->msgs_len; i++)
-    upb_decoder_jit_assignmsglabs(&h->msgs[i], &pclabel_count);
+    upb_decoder_jit_assignmsglabs(h, &h->msgs[i], &pclabel_count);
   for (int i = 0; i < h->msgs_len; i++)
     upb_decoder_jit_assignmsglabs2(h, &h->msgs[i]);
 
@@ -608,6 +625,9 @@ void upb_decoder_makejit(upb_decoder *d) {
         m->tablearray[j] = d->jit_code + dasm_getpclabel(d, 0);
       }
     }
+    if (m->endgroup_f) {
+      m->tablearray[m->endgroup_f->number] = d->jit_code + dasm_getpclabel(d, m->endgroup_f->jit_pclabel);
+    }
   }
 
   // Create debug info.
diff --git a/src/upb_msg.c b/src/upb_msg.c
index aac2c91..cf0625e 100644
--- a/src/upb_msg.c
+++ b/src/upb_msg.c
@@ -300,11 +300,11 @@ static upb_flow_t upb_dmsgsink_value(void *_m, upb_value fval, upb_value val) {
   upb_msg *m = _m;
   upb_fielddef *f = upb_value_getfielddef(fval);
   if (upb_isstring(f)) {
-    //fprintf(stderr, "dmsg_value!  this=%p f=%p name=" UPB_STRFMT ",
-    //        " UPB_STRFMT "  %p\n", m, f, UPB_STRARG(f->name), UPB_STRARG(val.val.str));
+    //fprintf(stderr, "dmsg_value!  this=%p f=%p name=" UPB_STRFMT ", "
+    //        UPB_STRFMT "\n", m, f, UPB_STRARG(f->name), UPB_STRARG(val.val.str));
   } else {
-    //fprintf(stderr, "dmsg_value!  this=%p f=%p name=" UPB_STRFMT ",
-    //        %llu\n", m, f, UPB_STRARG(f->name), val.val.uint64);
+    //fprintf(stderr, "dmsg_value!  this=%p f=%p name=" UPB_STRFMT ", %llu\n",
+    //        m, f, UPB_STRARG(f->name), val.val.uint64);
   }
   upb_msg_appendval(m, f, val);
   return UPB_CONTINUE;
@@ -320,12 +320,18 @@ static upb_sflow_t upb_dmsgsink_startsubmsg(void *_m, upb_value fval) {
   return UPB_CONTINUE_WITH(p);
 }
 
+static upb_flow_t endsubmsg(void *closure, upb_value fval) {
+  upb_fielddef *f = upb_value_getfielddef(fval);
+  //fprintf(stderr, "dmsg_endsubmsg!  " UPB_STRFMT " %p\n", UPB_STRARG(f->name), f);
+  return UPB_CONTINUE;
+}
+
 void upb_msg_regdhandlers(upb_handlers *h) {
   upb_register_all(h,
       NULL,  // startmsg
       NULL,  // endmsg
       &upb_dmsgsink_value,
       &upb_dmsgsink_startsubmsg,
-      NULL,  // endsubmsg
+      endsubmsg,  // endsubmsg
       NULL); // unknown
 }
diff --git a/src/upb_stream.c b/src/upb_stream.c
index 982c8a3..a408925 100644
--- a/src/upb_stream.c
+++ b/src/upb_stream.c
@@ -52,7 +52,7 @@ static void upb_msgent_init(upb_handlers_msgent *e) {
   e->startmsg = &upb_startmsg_nop;
   e->endmsg = &upb_endmsg_nop;
   e->unknownval = &upb_unknownval_nop;
-  e->is_group = false;
+  e->endgroup_f = NULL;
   e->tablearray = NULL;
 }
 
@@ -76,6 +76,7 @@ void upb_handlers_uninit(upb_handlers *h) {
   for (int i = 0; i < h->msgs_len; i++) {
     upb_inttable_free(&h->msgs[i].fieldtab);
     free(h->msgs[i].tablearray);
+    free(h->msgs[i].endgroup_f);
   }
   free(h->msgs);
   upb_msgdef_unref(h->toplevel_msgdef);
@@ -87,8 +88,8 @@ static upb_handlers_fieldent *upb_handlers_getorcreate_without_fval(
   upb_handlers_fieldent *f =
       upb_inttable_lookup(&h->msgent->fieldtab, tag);
   if (!f) {
-    upb_handlers_fieldent new_f = {false, type, -1, UPB_NO_VALUE,
-        {&upb_value_nop}, &upb_endsubmsg_nop, 0, 0, 0, repeated};
+    upb_handlers_fieldent new_f = {false, type, repeated, fieldnum, -1, UPB_NO_VALUE,
+        {&upb_value_nop}, &upb_endsubmsg_nop, 0, 0, 0};
     if (upb_issubmsgtype(type)) new_f.cb.startsubmsg = &upb_startsubmsg_nop;
     upb_inttable_insert(&h->msgent->fieldtab, tag, &new_f);
 
@@ -226,14 +227,14 @@ void upb_handlers_pop(upb_handlers *h, upb_fielddef *f) {
 /* upb_dispatcher *************************************************************/
 
 static upb_handlers_fieldent toplevel_f = {
-  false, UPB_TYPE(GROUP),
+  false, UPB_TYPE(GROUP), false, 0,
   0, // msgent_index
 #ifdef NDEBUG
   {{0}},
 #else
   {{0}, UPB_VALUETYPE_RAW},
 #endif
-  {NULL}, NULL, 0, 0, 0, false};
+  {NULL}, NULL, 0, 0, 0};
 
 void upb_dispatcher_init(upb_dispatcher *d, upb_handlers *h) {
   d->handlers = h;
diff --git a/src/upb_stream.h b/src/upb_stream.h
index 7ae9b8d..d024675 100644
--- a/src/upb_stream.h
+++ b/src/upb_stream.h
@@ -92,6 +92,8 @@ upb_flow_t upb_unknownval_nop(void *closure, upb_field_number_t fieldnum,
 typedef struct {
   bool junk;
   upb_fieldtype_t type;
+  bool repeated;
+  uint32_t number;
   // For upb_issubmsg(f) only, the index into the msgdef array of the submsg.
   // -1 if unset (indicates that submsg should be skipped).
   int32_t msgent_index;
@@ -104,7 +106,6 @@ typedef struct {
   uint32_t jit_pclabel;
   uint32_t jit_pclabel_notypecheck;
   uint32_t jit_submsg_done_pclabel;
-  bool repeated;
 } upb_handlers_fieldent;
 
 typedef struct _upb_handlers_msgent {
@@ -117,8 +118,7 @@ typedef struct _upb_handlers_msgent {
   uint32_t jit_endofbuf_pclabel;
   uint32_t jit_endofmsg_pclabel;
   uint32_t jit_unknownfield_pclabel;
-  uint32_t groupnum;
-  bool is_group;
+  upb_handlers_fieldent *endgroup_f;  // NULL if not a group.
   int32_t jit_parent_field_done_pclabel;
   uint32_t max_field_number;
   void **tablearray;