Google-internal development.

11 files changed, 474 insertions, 395 deletions

diff --git a/upb/pb/compile_decoder.c b/upb/pb/compile_decoder.c
index 11aa4e9..a17332b 100644
--- a/upb/pb/compile_decoder.c
+++ b/upb/pb/compile_decoder.c
@@ -64,7 +64,6 @@ mgroup *newgroup(const void *owner) {
 
 static void freemethod(upb_refcounted *r) {
   upb_pbdecodermethod *method = (upb_pbdecodermethod*)r;
-  upb_byteshandler_uninit(&method->input_handler_);
 
   if (method->dest_handlers_) {
     upb_handlers_unref(method->dest_handlers_, method);
@@ -762,8 +761,10 @@ static void compile_method(compiler *c, upb_pbdecodermethod *method) {
   putsel(c, OP_STARTMSG, UPB_STARTMSG_SELECTOR, h);
  label(c, LABEL_FIELD);
   uint32_t* start_pc = c->pc;
-  upb_msg_iter i;
-  for(upb_msg_begin(&i, md); !upb_msg_done(&i); upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for(upb_msg_field_begin(&i, md);
+      !upb_msg_field_done(&i);
+      upb_msg_field_next(&i)) {
     const upb_fielddef *f = upb_msg_iter_field(&i);
     upb_fieldtype_t type = upb_fielddef_type(f);
 
@@ -813,9 +814,11 @@ static void find_methods(compiler *c, const upb_handlers *h) {
   newmethod(h, c->group);
 
   // Find submethods.
-  upb_msg_iter i;
+  upb_msg_field_iter i;
   const upb_msgdef *md = upb_handlers_msgdef(h);
-  for(upb_msg_begin(&i, md); !upb_msg_done(&i); upb_msg_next(&i)) {
+  for(upb_msg_field_begin(&i, md);
+      !upb_msg_field_done(&i);
+      upb_msg_field_next(&i)) {
     const upb_fielddef *f = upb_msg_iter_field(&i);
     const upb_handlers *sub_h;
     if (upb_fielddef_type(f) == UPB_TYPE_MESSAGE &&
@@ -857,7 +860,7 @@ static void set_bytecode_handlers(mgroup *g) {
 }
 
 
-/* JIT setup. ******************************************************************/
+/* JIT setup. *****************************************************************/
 
 #ifdef UPB_USE_JIT_X64
 
diff --git a/upb/pb/compile_decoder_x64.dasc b/upb/pb/compile_decoder_x64.dasc
index 3181cab..e72e4e3 100644
--- a/upb/pb/compile_decoder_x64.dasc
+++ b/upb/pb/compile_decoder_x64.dasc
@@ -242,7 +242,7 @@ static void emit_static_asm(jitcompiler *jc) {
   |
   |2:
   |  // Resume decoder.
-  |  lea   ARG2_64, DECODER->callstack
+  |  mov   ARG2_64, DECODER->callstack
   |  sub   rsp, ARG3_64
   |  mov   ARG1_64, rsp
   |  callp memcpy  // Restore stack.
@@ -255,7 +255,7 @@ static void emit_static_asm(jitcompiler *jc) {
   asmlabel(jc, "exitjit");
   |->exitjit:
   |  // Save the stack into DECODER->callstack.
-  |  lea   ARG1_64, DECODER->callstack
+  |  mov   ARG1_64, DECODER->callstack
   |  mov   ARG2_64, rsp
   |  mov   ARG3_64, DECODER->saved_rsp
   |  sub   ARG3_64, rsp
@@ -300,11 +300,11 @@ static void emit_static_asm(jitcompiler *jc) {
   |  sub   rcx, rdx
   |  jb    ->err  // Len is greater than enclosing message.
   |  mov   FRAME->end_ofs, rcx
+  |  cmp   FRAME, DECODER->limit
+  |  je    >3   // Stack overflow
   |  add   FRAME, sizeof(upb_pbdecoder_frame)
   |  mov   DELIMEND, PTR
   |  add   DELIMEND, rdx
-  |  cmp   FRAME, DECODER->limit
-  |  je    >3   // Stack overflow
   |  mov   dword FRAME->groupnum, 0
   |  test  rcx, rcx
   |  jz    >2
@@ -1071,9 +1071,9 @@ static void jitbytecode(jitcompiler *jc) {
       |  // code with the packed code-path.  If this is changed later, this
       |  // store can be removed.
       |  mov   qword FRAME->end_ofs, 0
-      |  add   FRAME, sizeof(upb_pbdecoder_frame)
       |  cmp   FRAME, DECODER->limit
       |  je    ->err
+      |  add   FRAME, sizeof(upb_pbdecoder_frame)
       |  mov   dword FRAME->groupnum, arg
       break;
     case OP_PUSHLENDELIM:
diff --git a/upb/pb/compile_decoder_x64.h b/upb/pb/compile_decoder_x64.h
index ef4459d..9527361 100644
--- a/upb/pb/compile_decoder_x64.h
+++ b/upb/pb/compile_decoder_x64.h
@@ -28,8 +28,8 @@ static const unsigned char upb_jit_actionlist[2162] = {
   73,139,159,233,77,139,167,233,77,139,174,233,73,139,174,233,73,43,175,233,
   73,3,175,233,73,139,151,233,72,133,210,15,133,244,248,252,255,208,73,139,
   135,233,73,199,135,233,0,0,0,0,248,1,255,91,65,92,65,93,65,94,65,95,93,195,
-  248,2,73,141,183,233,72,41,212,72,137,231,72,184,237,237,65,84,73,137,228,
-  72,129,228,239,252,255,208,76,137,228,65,92,195,255,248,11,73,141,191,233,
+  248,2,73,139,183,233,72,41,212,72,137,231,72,184,237,237,65,84,73,137,228,
+  72,129,228,239,252,255,208,76,137,228,65,92,195,255,248,11,73,139,191,233,
   72,137,230,73,139,151,233,72,41,226,73,137,151,233,137,195,72,184,237,237,
   65,84,73,137,228,72,129,228,239,252,255,208,76,137,228,65,92,137,216,73,139,
   167,233,91,65,92,65,93,65,94,65,95,93,195,255,248,12,73,57,159,233,15,132,
@@ -40,7 +40,7 @@ static const unsigned char upb_jit_actionlist[2162] = {
   255,76,57,227,15,132,244,253,255,76,137,225,72,41,217,72,131,252,249,1,15,
   130,244,253,255,15,182,19,132,210,15,137,244,254,248,7,232,244,14,248,8,72,
   131,195,1,72,137,252,233,72,41,217,72,41,209,15,130,244,15,73,137,142,233,
-  73,129,198,239,72,137,221,72,1,213,77,59,183,233,15,132,244,249,65,199,134,
+  77,59,183,233,15,132,244,249,73,129,198,239,72,137,221,72,1,213,65,199,134,
   233,0,0,0,0,72,133,201,15,132,244,248,77,139,167,233,72,57,252,235,15,135,
   244,248,76,57,229,15,135,244,248,255,73,137,252,236,248,2,195,248,3,73,139,
   159,233,76,137,252,255,255,72,190,237,237,255,190,237,255,49,252,246,255,
@@ -122,8 +122,8 @@ static const unsigned char upb_jit_actionlist[2162] = {
   1,248,2,255,72,137,218,76,137,225,72,41,217,77,139,135,233,72,184,237,237,
   65,84,73,137,228,72,129,228,239,252,255,208,76,137,228,65,92,72,1,195,255,
   76,57,227,15,132,244,249,232,244,29,248,3,255,76,137,227,255,72,57,252,235,
-  15,133,244,1,248,4,255,77,137,174,233,73,199,134,233,0,0,0,0,73,129,198,239,
-  77,59,183,233,15,132,244,15,65,199,134,233,237,255,232,244,13,255,73,129,
+  15,133,244,1,248,4,255,77,137,174,233,73,199,134,233,0,0,0,0,77,59,183,233,
+  15,132,244,15,73,129,198,239,65,199,134,233,237,255,232,244,13,255,73,129,
   252,238,239,77,139,174,233,255,77,139,167,233,73,3,174,233,73,59,175,233,
   15,130,244,247,76,57,229,15,135,244,247,73,137,252,236,248,1,255,72,57,221,
   15,132,245,255,232,245,255,248,9,72,131,196,8,195,255
@@ -419,7 +419,7 @@ static void emit_static_asm(jitcompiler *jc) {
   //|
   //|2:
   //|  // Resume decoder.
-  //|  lea   ARG2_64, DECODER->callstack
+  //|  mov   ARG2_64, DECODER->callstack
   //|  sub   rsp, ARG3_64
   //|  mov   ARG1_64, rsp
   //|  callp memcpy  // Restore stack.
@@ -434,7 +434,7 @@ static void emit_static_asm(jitcompiler *jc) {
   asmlabel(jc, "exitjit");
   //|->exitjit:
   //|  // Save the stack into DECODER->callstack.
-  //|  lea   ARG1_64, DECODER->callstack
+  //|  mov   ARG1_64, DECODER->callstack
   //|  mov   ARG2_64, rsp
   //|  mov   ARG3_64, DECODER->saved_rsp
   //|  sub   ARG3_64, rsp
@@ -490,11 +490,11 @@ static void emit_static_asm(jitcompiler *jc) {
   //|  sub   rcx, rdx
   //|  jb    ->err  // Len is greater than enclosing message.
   //|  mov   FRAME->end_ofs, rcx
+  //|  cmp   FRAME, DECODER->limit
+  //|  je    >3   // Stack overflow
   //|  add   FRAME, sizeof(upb_pbdecoder_frame)
   //|  mov   DELIMEND, PTR
   //|  add   DELIMEND, rdx
-  //|  cmp   FRAME, DECODER->limit
-  //|  je    >3   // Stack overflow
   //|  mov   dword FRAME->groupnum, 0
   //|  test  rcx, rcx
   //|  jz    >2
@@ -504,7 +504,7 @@ static void emit_static_asm(jitcompiler *jc) {
   //|  cmp   DELIMEND, DATAEND
   //|  ja    >2
   //|  mov   DATAEND, DELIMEND  // If DELIMEND >= PTR && DELIMEND < DATAEND
-  dasm_put(Dst, 337, Dt1(->end_ofs), sizeof(upb_pbdecoder_frame), Dt2(->limit), Dt1(->groupnum), Dt2(->end));
+  dasm_put(Dst, 337, Dt1(->end_ofs), Dt2(->limit), sizeof(upb_pbdecoder_frame), Dt1(->groupnum), Dt2(->end));
 # 317 "upb/pb/compile_decoder_x64.dasc"
   //|2:
   //|  ret
@@ -1609,11 +1609,11 @@ static void jitbytecode(jitcompiler *jc) {
       //|  // code with the packed code-path.  If this is changed later, this
       //|  // store can be removed.
       //|  mov   qword FRAME->end_ofs, 0
-      //|  add   FRAME, sizeof(upb_pbdecoder_frame)
       //|  cmp   FRAME, DECODER->limit
       //|  je    ->err
+      //|  add   FRAME, sizeof(upb_pbdecoder_frame)
       //|  mov   dword FRAME->groupnum, arg
-      dasm_put(Dst, 2070, Dt1(->sink.closure), Dt1(->end_ofs), sizeof(upb_pbdecoder_frame), Dt2(->limit), Dt1(->groupnum), arg);
+      dasm_put(Dst, 2070, Dt1(->sink.closure), Dt1(->end_ofs), Dt2(->limit), sizeof(upb_pbdecoder_frame), Dt1(->groupnum), arg);
 # 1078 "upb/pb/compile_decoder_x64.dasc"
       break;
     case OP_PUSHLENDELIM:
diff --git a/upb/pb/decoder.c b/upb/pb/decoder.c
index 04ca413..f371bce 100644
--- a/upb/pb/decoder.c
+++ b/upb/pb/decoder.c
@@ -19,10 +19,7 @@
  */
 
 #include <inttypes.h>
-#include <setjmp.h>
-#include <stdarg.h>
 #include <stddef.h>
-#include <stdlib.h>
 #include "upb/pb/decoder.int.h"
 #include "upb/pb/varint.int.h"
 
@@ -70,18 +67,17 @@ static bool consumes_input(opcode op) {
 
 static bool in_residual_buf(const upb_pbdecoder *d, const char *p);
 
-// It's unfortunate that we have to micro-manage the compiler this way,
-// especially since this tuning is necessarily specific to one hardware
-// configuration.  But emperically on a Core i7, performance increases 30-50%
-// with these annotations.  Every instance where these appear, gcc 4.2.1 made
-// the wrong decision and degraded performance in benchmarks.
-#define FORCEINLINE static inline __attribute__((always_inline))
-#define NOINLINE __attribute__((noinline))
+// It's unfortunate that we have to micro-manage the compiler with
+// UPB_FORCEINLINE and UPB_NOINLINE, especially since this tuning is necessarily
+// specific to one hardware configuration.  But empirically on a Core i7,
+// performance increases 30-50% with these annotations.  Every instance where
+// these appear, gcc 4.2.1 made the wrong decision and degraded performance in
+// benchmarks.
 
 static void seterr(upb_pbdecoder *d, const char *msg) {
-  // TODO(haberman): encapsulate this access to pipeline->status, but not sure
-  // exactly what that interface should look like.
-  upb_status_seterrmsg(d->status, msg);
+  upb_status status = UPB_STATUS_INIT;
+  upb_status_seterrmsg(&status, msg);
+  upb_env_reporterror(d->env, &status);
 }
 
 void upb_pbdecoder_seterr(upb_pbdecoder *d, const char *msg) {
@@ -249,7 +245,8 @@ static int32_t skip(upb_pbdecoder *d, size_t bytes) {
 
 // Copies the next "bytes" bytes into "buf" and advances the stream.
 // Requires that this many bytes are available in the current buffer.
-FORCEINLINE void consumebytes(upb_pbdecoder *d, void *buf, size_t bytes) {
+UPB_FORCEINLINE static void consumebytes(upb_pbdecoder *d, void *buf,
+                                         size_t bytes) {
   assert(bytes <= curbufleft(d));
   memcpy(buf, d->ptr, bytes);
   advance(d, bytes);
@@ -258,8 +255,8 @@ FORCEINLINE void consumebytes(upb_pbdecoder *d, void *buf, size_t bytes) {
 // Slow path for getting the next "bytes" bytes, regardless of whether they are
 // available in the current buffer or not.  Returns a status code as described
 // in decoder.int.h.
-static NOINLINE int32_t getbytes_slow(upb_pbdecoder *d, void *buf,
-                                      size_t bytes) {
+UPB_NOINLINE static int32_t getbytes_slow(upb_pbdecoder *d, void *buf,
+                                          size_t bytes) {
   const size_t avail = curbufleft(d);
   consumebytes(d, buf, avail);
   bytes -= avail;
@@ -280,7 +277,8 @@ static NOINLINE int32_t getbytes_slow(upb_pbdecoder *d, void *buf,
 
 // Gets the next "bytes" bytes, regardless of whether they are available in the
 // current buffer or not.  Returns a status code as described in decoder.int.h.
-FORCEINLINE int32_t getbytes(upb_pbdecoder *d, void *buf, size_t bytes) {
+UPB_FORCEINLINE static int32_t getbytes(upb_pbdecoder *d, void *buf,
+                                        size_t bytes) {
   if (curbufleft(d) >= bytes) {
     // Buffer has enough data to satisfy.
     consumebytes(d, buf, bytes);
@@ -290,8 +288,8 @@ FORCEINLINE int32_t getbytes(upb_pbdecoder *d, void *buf, size_t bytes) {
   }
 }
 
-static NOINLINE size_t peekbytes_slow(upb_pbdecoder *d, void *buf,
-                                      size_t bytes) {
+UPB_NOINLINE static size_t peekbytes_slow(upb_pbdecoder *d, void *buf,
+                                          size_t bytes) {
   size_t ret = curbufleft(d);
   memcpy(buf, d->ptr, ret);
   if (in_residual_buf(d, d->ptr)) {
@@ -302,7 +300,8 @@ static NOINLINE size_t peekbytes_slow(upb_pbdecoder *d, void *buf,
   return ret;
 }
 
-FORCEINLINE size_t peekbytes(upb_pbdecoder *d, void *buf, size_t bytes) {
+UPB_FORCEINLINE static size_t peekbytes(upb_pbdecoder *d, void *buf,
+                                        size_t bytes) {
   if (curbufleft(d) >= bytes) {
     memcpy(buf, d->ptr, bytes);
     return bytes;
@@ -316,8 +315,8 @@ FORCEINLINE size_t peekbytes(upb_pbdecoder *d, void *buf, size_t bytes) {
 
 // Slow path for decoding a varint from the current buffer position.
 // Returns a status code as described in decoder.int.h.
-NOINLINE int32_t upb_pbdecoder_decode_varint_slow(upb_pbdecoder *d,
-                                                  uint64_t *u64) {
+UPB_NOINLINE int32_t upb_pbdecoder_decode_varint_slow(upb_pbdecoder *d,
+                                                      uint64_t *u64) {
   *u64 = 0;
   uint8_t byte = 0x80;
   int bitpos;
@@ -335,7 +334,7 @@ NOINLINE int32_t upb_pbdecoder_decode_varint_slow(upb_pbdecoder *d,
 
 // Decodes a varint from the current buffer position.
 // Returns a status code as described in decoder.int.h.
-FORCEINLINE int32_t decode_varint(upb_pbdecoder *d, uint64_t *u64) {
+UPB_FORCEINLINE static int32_t decode_varint(upb_pbdecoder *d, uint64_t *u64) {
   if (curbufleft(d) > 0 && !(*d->ptr & 0x80)) {
     *u64 = *d->ptr;
     advance(d, 1);
@@ -358,7 +357,7 @@ FORCEINLINE int32_t decode_varint(upb_pbdecoder *d, uint64_t *u64) {
 
 // Decodes a 32-bit varint from the current buffer position.
 // Returns a status code as described in decoder.int.h.
-FORCEINLINE int32_t decode_v32(upb_pbdecoder *d, uint32_t *u32) {
+UPB_FORCEINLINE static int32_t decode_v32(upb_pbdecoder *d, uint32_t *u32) {
   uint64_t u64;
   int32_t ret = decode_varint(d, &u64);
   if (ret >= 0) return ret;
@@ -377,14 +376,14 @@ FORCEINLINE int32_t decode_v32(upb_pbdecoder *d, uint32_t *u32) {
 // Decodes a fixed32 from the current buffer position.
 // Returns a status code as described in decoder.int.h.
 // TODO: proper byte swapping for big-endian machines.
-FORCEINLINE int32_t decode_fixed32(upb_pbdecoder *d, uint32_t *u32) {
+UPB_FORCEINLINE static int32_t decode_fixed32(upb_pbdecoder *d, uint32_t *u32) {
   return getbytes(d, u32, 4);
 }
 
 // Decodes a fixed64 from the current buffer position.
 // Returns a status code as described in decoder.int.h.
 // TODO: proper byte swapping for big-endian machines.
-FORCEINLINE int32_t decode_fixed64(upb_pbdecoder *d, uint64_t *u64) {
+UPB_FORCEINLINE static int32_t decode_fixed64(upb_pbdecoder *d, uint64_t *u64) {
   return getbytes(d, u64, 8);
 }
 
@@ -408,7 +407,7 @@ static bool decoder_push(upb_pbdecoder *d, uint64_t end) {
   if (end > fr->end_ofs) {
     seterr(d, "Submessage end extends past enclosing submessage.");
     return false;
-  } else if ((fr + 1) == d->limit) {
+  } else if (fr == d->limit) {
     seterr(d, kPbDecoderStackOverflow);
     return false;
   }
@@ -435,8 +434,8 @@ static bool pushtagdelim(upb_pbdecoder *d, uint32_t arg) {
 // Pops a frame from the decoder stack.
 static void decoder_pop(upb_pbdecoder *d) { d->top--; }
 
-NOINLINE int32_t upb_pbdecoder_checktag_slow(upb_pbdecoder *d,
-                                             uint64_t expected) {
+UPB_NOINLINE int32_t upb_pbdecoder_checktag_slow(upb_pbdecoder *d,
+                                                 uint64_t expected) {
   uint64_t data = 0;
   size_t bytes = upb_value_size(expected);
   size_t read = peekbytes(d, &data, bytes);
@@ -814,7 +813,10 @@ size_t upb_pbdecoder_decode(void *closure, const void *hd, const char *buf,
 void *upb_pbdecoder_startbc(void *closure, const void *pc, size_t size_hint) {
   upb_pbdecoder *d = closure;
   UPB_UNUSED(size_hint);
+  d->top->end_ofs = UINT64_MAX;
+  d->bufstart_ofs = 0;
   d->call_len = 1;
+  d->callstack[0] = &halt;
   d->pc = pc;
   return d;
 }
@@ -823,6 +825,8 @@ void *upb_pbdecoder_startjit(void *closure, const void *hd, size_t size_hint) {
   UPB_UNUSED(hd);
   UPB_UNUSED(size_hint);
   upb_pbdecoder *d = closure;
+  d->top->end_ofs = UINT64_MAX;
+  d->bufstart_ofs = 0;
   d->call_len = 0;
   return d;
 }
@@ -879,55 +883,115 @@ bool upb_pbdecoder_end(void *closure, const void *handler_data) {
   return true;
 }
 
-void upb_pbdecoder_init(upb_pbdecoder *d, const upb_pbdecodermethod *m,
-                        upb_status *s) {
-  d->limit = &d->stack[UPB_DECODER_MAX_NESTING];
-  upb_bytessink_reset(&d->input_, &m->input_handler_, d);
-  d->method_ = m;
-  d->callstack[0] = &halt;
-  d->status = s;
-  upb_pbdecoder_reset(d);
-}
-
 void upb_pbdecoder_reset(upb_pbdecoder *d) {
   d->top = d->stack;
-  d->top->end_ofs = UINT64_MAX;
   d->top->groupnum = 0;
-  d->bufstart_ofs = 0;
   d->ptr = d->residual;
   d->buf = d->residual;
   d->end = d->residual;
   d->residual_end = d->residual;
-  d->call_len = 1;
 }
 
-uint64_t upb_pbdecoder_bytesparsed(const upb_pbdecoder *d) {
-  return offset(d);
+static size_t stacksize(upb_pbdecoder *d, size_t entries) {
+  UPB_UNUSED(d);
+  return entries * sizeof(upb_pbdecoder_frame);
 }
 
-// Not currently required, but to support outgrowing the static stack we need
-// this.
-void upb_pbdecoder_uninit(upb_pbdecoder *d) {
+static size_t callstacksize(upb_pbdecoder *d, size_t entries) {
   UPB_UNUSED(d);
-}
 
-const upb_pbdecodermethod *upb_pbdecoder_method(const upb_pbdecoder *d) {
-  return d->method_;
+#ifdef UPB_USE_JIT_X64
+  if (d->method_->is_native_) {
+    // Each native stack frame needs two pointers, plus we need a few frames for
+    // the enter/exit trampolines.
+    size_t ret = entries * sizeof(void*) * 2;
+    ret += sizeof(void*) * 10;
+    return ret;
+  }
+#endif
+
+  return entries * sizeof(uint32_t*);
 }
 
-bool upb_pbdecoder_resetoutput(upb_pbdecoder *d, upb_sink* sink) {
-  // TODO(haberman): do we need to test whether the decoder is already on the
-  // stack (like calling this from within a callback)?  Should we support
-  // rebinding the output at all?
+upb_pbdecoder *upb_pbdecoder_create(upb_env *e, const upb_pbdecodermethod *m,
+                                    upb_sink *sink) {
+  const size_t default_max_nesting = 64;
+#ifndef NDEBUG
+  size_t size_before = upb_env_bytesallocated(e);
+#endif
+
+  upb_pbdecoder *d = upb_env_malloc(e, sizeof(upb_pbdecoder));
+  if (!d) return NULL;
+
+  d->method_ = m;
+  d->callstack = upb_env_malloc(e, callstacksize(d, default_max_nesting));
+  d->stack = upb_env_malloc(e, stacksize(d, default_max_nesting));
+  if (!d->stack || !d->callstack) {
+    return NULL;
+  }
+
+  d->env = e;
+  d->limit = d->stack + default_max_nesting - 1;
+  d->stack_size = default_max_nesting;
+
+  upb_pbdecoder_reset(d);
+  upb_bytessink_reset(&d->input_, &m->input_handler_, d);
+
   assert(sink);
   if (d->method_->dest_handlers_) {
     if (sink->handlers != d->method_->dest_handlers_)
-      return false;
+      return NULL;
   }
   upb_sink_reset(&d->top->sink, sink->handlers, sink->closure);
-  return true;
+
+  // If this fails, increase the value in decoder.h.
+  assert(upb_env_bytesallocated(e) - size_before <= UPB_PB_DECODER_SIZE);
+  return d;
+}
+
+uint64_t upb_pbdecoder_bytesparsed(const upb_pbdecoder *d) {
+  return offset(d);
+}
+
+const upb_pbdecodermethod *upb_pbdecoder_method(const upb_pbdecoder *d) {
+  return d->method_;
 }
 
 upb_bytessink *upb_pbdecoder_input(upb_pbdecoder *d) {
   return &d->input_;
 }
+
+size_t upb_pbdecoder_maxnesting(const upb_pbdecoder *d) {
+  return d->stack_size;
+}
+
+bool upb_pbdecoder_setmaxnesting(upb_pbdecoder *d, size_t max) {
+  if (max < d->top - d->stack) {
+    // Can't set a limit smaller than what we are currently at.
+    return false;
+  }
+
+  if (max > d->stack_size) {
+    // Need to reallocate stack and callstack to accommodate.
+    size_t old_size = stacksize(d, d->stack_size);
+    size_t new_size = stacksize(d, max);
+    void *p = upb_env_realloc(d->env, d->stack, old_size, new_size);
+    if (!p) {
+      return false;
+    }
+    d->stack = p;
+
+    old_size = callstacksize(d, d->stack_size);
+    new_size = callstacksize(d, max);
+    p = upb_env_realloc(d->env, d->callstack, old_size, new_size);
+    if (!p) {
+      return false;
+    }
+    d->callstack = p;
+
+    d->stack_size = max;
+  }
+
+  d->limit = d->stack + max - 1;
+  return true;
+}
diff --git a/upb/pb/decoder.h b/upb/pb/decoder.h
index e93f3bd..d37718c 100644
--- a/upb/pb/decoder.h
+++ b/upb/pb/decoder.h
@@ -18,7 +18,7 @@
 #ifndef UPB_DECODER_H_
 #define UPB_DECODER_H_
 
-#include "upb/table.int.h"
+#include "upb/env.h"
 #include "upb/sink.h"
 
 #ifdef __cplusplus
@@ -37,44 +37,6 @@ UPB_DECLARE_TYPE(upb::pb::Decoder, upb_pbdecoder);
 UPB_DECLARE_TYPE(upb::pb::DecoderMethod, upb_pbdecodermethod);
 UPB_DECLARE_TYPE(upb::pb::DecoderMethodOptions, upb_pbdecodermethodopts);
 
-// The maximum that any submessages can be nested.  Matches proto2's limit.
-// This specifies the size of the decoder's statically-sized array and therefore
-// setting it high will cause the upb::pb::Decoder object to be larger.
-//
-// If necessary we can add a runtime-settable property to Decoder that allow
-// this to be larger than the compile-time setting, but this would add
-// complexity, particularly since we would have to decide how/if to give users
-// the ability to set a custom memory allocation function.
-#define UPB_DECODER_MAX_NESTING 64
-
-// Internal-only struct used by the decoder.
-typedef struct {
- UPB_PRIVATE_FOR_CPP
-  // Space optimization note: we store two pointers here that the JIT
-  // doesn't need at all; the upb_handlers* inside the sink and
-  // the dispatch table pointer.  We can optimze so that the JIT uses
-  // smaller stack frames than the interpreter.  The only thing we need
-  // to guarantee is that the fallback routines can find end_ofs.
-  upb_sink sink;
-
-  // The absolute stream offset of the end-of-frame delimiter.
-  // Non-delimited frames (groups and non-packed repeated fields) reuse the
-  // delimiter of their parent, even though the frame may not end there.
-  //
-  // NOTE: the JIT stores a slightly different value here for non-top frames.
-  // It stores the value relative to the end of the enclosed message.  But the
-  // top frame is still stored the same way, which is important for ensuring
-  // that calls from the JIT into C work correctly.
-  uint64_t end_ofs;
-  const uint32_t *base;
-
-  // 0 indicates a length-delimited field.
-  // A positive number indicates a known group.
-  // A negative number indicates an unknown group.
-  int32_t groupnum;
-  upb_inttable *dispatch;  // Not used by the JIT.
-} upb_pbdecoder_frame;
-
 // The parameters one uses to construct a DecoderMethod.
 // TODO(haberman): move allowjit here?  Seems more convenient for users.
 UPB_DEFINE_CLASS0(upb::pb::DecoderMethodOptions,
@@ -152,22 +114,31 @@ UPB_DEFINE_STRUCT(upb_pbdecodermethod, upb_refcounted,
   upb_inttable dispatch;
 ));
 
+// Preallocation hint: decoder won't allocate more bytes than this when first
+// constructed.  This hint may be an overestimate for some build configurations.
+// But if the decoder library is upgraded without recompiling the application,
+// it may be an underestimate.
+#define UPB_PB_DECODER_SIZE 4400
+
+#ifdef __cplusplus
+
 // A Decoder receives binary protobuf data on its input sink and pushes the
 // decoded data to its output sink.
-UPB_DEFINE_CLASS0(upb::pb::Decoder,
+class upb::pb::Decoder {
  public:
   // Constructs a decoder instance for the given method, which must outlive this
   // decoder.  Any errors during parsing will be set on the given status, which
   // must also outlive this decoder.
-  Decoder(const DecoderMethod* method, Status* status);
-  ~Decoder();
+  //
+  // The sink must match the given method.
+  static Decoder* Create(Environment* env, const DecoderMethod* method,
+                         Sink* output);
 
   // Returns the DecoderMethod this decoder is parsing from.
-  // TODO(haberman): Do users need to be able to rebind this?
   const DecoderMethod* method() const;
 
-  // Resets the state of the decoder.
-  void Reset();
+  // The sink on which this decoder receives input.
+  BytesSink* input();
 
   // Returns number of bytes successfully parsed.
   //
@@ -178,76 +149,25 @@ UPB_DEFINE_CLASS0(upb::pb::Decoder,
   // callback.
   uint64_t BytesParsed() const;
 
-  // Resets the output sink of the Decoder.
-  // The given sink must match method()->dest_handlers().
+  // Gets/sets the parsing nexting limit.  If the total number of nested
+  // submessages and repeated fields hits this limit, parsing will fail.  This
+  // is a resource limit that controls the amount of memory used by the parsing
+  // stack.
   //
-  // This must be called at least once before the decoder can be used.  It may
-  // only be called with the decoder is in a state where it was just created or
-  // reset with pipeline.Reset().  The given sink must be from the same pipeline
-  // as this decoder.
-  bool ResetOutput(Sink* sink);
-
-  // The sink on which this decoder receives input.
-  BytesSink* input();
-
- private:
-  UPB_DISALLOW_COPY_AND_ASSIGN(Decoder);
-,
-UPB_DEFINE_STRUCT0(upb_pbdecoder, UPB_QUOTE(
-  // Our input sink.
-  upb_bytessink input_;
-
-  // The decoder method we are parsing with (owned).
-  const upb_pbdecodermethod *method_;
-
-  size_t call_len;
-  const uint32_t *pc, *last;
+  // Setting the limit will fail if the parser is currently suspended at a depth
+  // greater than this, or if memory allocation of the stack fails.
+  size_t max_nesting() const;
+  bool set_max_nesting(size_t max);
 
-  // Current input buffer and its stream offset.
-  const char *buf, *ptr, *end, *checkpoint;
-
-  // End of the delimited region, relative to ptr, or NULL if not in this buf.
-  const char *delim_end;
-
-  // End of the delimited region, relative to ptr, or end if not in this buf.
-  const char *data_end;
-
-  // Overall stream offset of "buf."
-  uint64_t bufstart_ofs;
-
-  // Buffer for residual bytes not parsed from the previous buffer.
-  // The maximum number of residual bytes we require is 12; a five-byte
-  // unknown tag plus an eight-byte value, less one because the value
-  // is only a partial value.
-  char residual[12];
-  char *residual_end;
-
-  // Stores the user buffer passed to our decode function.
-  const char *buf_param;
-  size_t size_param;
-  const upb_bufhandle *handle;
+  void Reset();
 
-#ifdef UPB_USE_JIT_X64
-  // Used momentarily by the generated code to store a value while a user
-  // function is called.
-  uint32_t tmp_len;
+  static const size_t kSize = UPB_PB_DECODER_SIZE;
 
-  const void *saved_rsp;
-#endif
+ private:
+  UPB_DISALLOW_POD_OPS(Decoder, upb::pb::Decoder);
+};
 
-  upb_status *status;
-
-  // Our internal stack.
-  upb_pbdecoder_frame *top, *limit;
-  upb_pbdecoder_frame stack[UPB_DECODER_MAX_NESTING];
-#ifdef UPB_USE_JIT_X64
-  // Each native stack frame needs two pointers, plus we need a few frames for
-  // the enter/exit trampolines.
-  const uint32_t *callstack[(UPB_DECODER_MAX_NESTING * 2) + 10];
-#else
-  const uint32_t *callstack[UPB_DECODER_MAX_NESTING];
-#endif
-)));
+#endif  // __cplusplus
 
 // A class for caching protobuf processing code, whether bytecode for the
 // interpreted decoder or machine code for the JIT.
@@ -296,14 +216,15 @@ UPB_DEFINE_STRUCT0(upb_pbcodecache,
 
 UPB_BEGIN_EXTERN_C  // {
 
-void upb_pbdecoder_init(upb_pbdecoder *d, const upb_pbdecodermethod *method,
-                        upb_status *status);
-void upb_pbdecoder_uninit(upb_pbdecoder *d);
-void upb_pbdecoder_reset(upb_pbdecoder *d);
+upb_pbdecoder *upb_pbdecoder_create(upb_env *e,
+                                    const upb_pbdecodermethod *method,
+                                    upb_sink *output);
 const upb_pbdecodermethod *upb_pbdecoder_method(const upb_pbdecoder *d);
-bool upb_pbdecoder_resetoutput(upb_pbdecoder *d, upb_sink *sink);
 upb_bytessink *upb_pbdecoder_input(upb_pbdecoder *d);
 uint64_t upb_pbdecoder_bytesparsed(const upb_pbdecoder *d);
+size_t upb_pbdecoder_maxnesting(const upb_pbdecoder *d);
+bool upb_pbdecoder_setmaxnesting(upb_pbdecoder *d, size_t max);
+void upb_pbdecoder_reset(upb_pbdecoder *d);
 
 void upb_pbdecodermethodopts_init(upb_pbdecodermethodopts *opts,
                                   const upb_handlers *h);
@@ -338,27 +259,27 @@ namespace upb {
 
 namespace pb {
 
-inline Decoder::Decoder(const DecoderMethod* m, Status* s) {
-  upb_pbdecoder_init(this, m, s);
-}
-inline Decoder::~Decoder() {
-  upb_pbdecoder_uninit(this);
+// static
+inline Decoder* Decoder::Create(Environment* env, const DecoderMethod* m,
+                                Sink* sink) {
+  return upb_pbdecoder_create(env, m, sink);
 }
 inline const DecoderMethod* Decoder::method() const {
   return upb_pbdecoder_method(this);
 }
-inline void Decoder::Reset() {
-  upb_pbdecoder_reset(this);
+inline BytesSink* Decoder::input() {
+  return upb_pbdecoder_input(this);
 }
 inline uint64_t Decoder::BytesParsed() const {
   return upb_pbdecoder_bytesparsed(this);
 }
-inline bool Decoder::ResetOutput(Sink* sink) {
-  return upb_pbdecoder_resetoutput(this, sink);
+inline size_t Decoder::max_nesting() const {
+  return upb_pbdecoder_maxnesting(this);
 }
-inline BytesSink* Decoder::input() {
-  return upb_pbdecoder_input(this);
+inline bool Decoder::set_max_nesting(size_t max) {
+  return upb_pbdecoder_setmaxnesting(this, max);
 }
+inline void Decoder::Reset() { upb_pbdecoder_reset(this); }
 
 inline DecoderMethodOptions::DecoderMethodOptions(const Handlers* h) {
   upb_pbdecodermethodopts_init(this, h);
diff --git a/upb/pb/decoder.int.h b/upb/pb/decoder.int.h
index 302701e..5522be7 100644
--- a/upb/pb/decoder.int.h
+++ b/upb/pb/decoder.int.h
@@ -13,8 +13,9 @@
 #include <stdlib.h>
 #include "upb/def.h"
 #include "upb/handlers.h"
-#include "upb/sink.h"
 #include "upb/pb/decoder.h"
+#include "upb/sink.h"
+#include "upb/table.int.h"
 
 // Opcode definitions.  The canonical meaning of each opcode is its
 // implementation in the interpreter (the JIT is written to match this).
@@ -112,6 +113,95 @@ typedef struct {
 #endif
 } mgroup;
 
+// The maximum that any submessages can be nested.  Matches proto2's limit.
+// This specifies the size of the decoder's statically-sized array and therefore
+// setting it high will cause the upb::pb::Decoder object to be larger.
+//
+// If necessary we can add a runtime-settable property to Decoder that allow
+// this to be larger than the compile-time setting, but this would add
+// complexity, particularly since we would have to decide how/if to give users
+// the ability to set a custom memory allocation function.
+#define UPB_DECODER_MAX_NESTING 64
+
+// Internal-only struct used by the decoder.
+typedef struct {
+  // Space optimization note: we store two pointers here that the JIT
+  // doesn't need at all; the upb_handlers* inside the sink and
+  // the dispatch table pointer.  We can optimze so that the JIT uses
+  // smaller stack frames than the interpreter.  The only thing we need
+  // to guarantee is that the fallback routines can find end_ofs.
+  upb_sink sink;
+
+  // The absolute stream offset of the end-of-frame delimiter.
+  // Non-delimited frames (groups and non-packed repeated fields) reuse the
+  // delimiter of their parent, even though the frame may not end there.
+  //
+  // NOTE: the JIT stores a slightly different value here for non-top frames.
+  // It stores the value relative to the end of the enclosed message.  But the
+  // top frame is still stored the same way, which is important for ensuring
+  // that calls from the JIT into C work correctly.
+  uint64_t end_ofs;
+  const uint32_t *base;
+
+  // 0 indicates a length-delimited field.
+  // A positive number indicates a known group.
+  // A negative number indicates an unknown group.
+  int32_t groupnum;
+  upb_inttable *dispatch;  // Not used by the JIT.
+} upb_pbdecoder_frame;
+
+struct upb_pbdecoder {
+  upb_env *env;
+
+  // Our input sink.
+  upb_bytessink input_;
+
+  // The decoder method we are parsing with (owned).
+  const upb_pbdecodermethod *method_;
+
+  size_t call_len;
+  const uint32_t *pc, *last;
+
+  // Current input buffer and its stream offset.
+  const char *buf, *ptr, *end, *checkpoint;
+
+  // End of the delimited region, relative to ptr, or NULL if not in this buf.
+  const char *delim_end;
+
+  // End of the delimited region, relative to ptr, or end if not in this buf.
+  const char *data_end;
+
+  // Overall stream offset of "buf."
+  uint64_t bufstart_ofs;
+
+  // Buffer for residual bytes not parsed from the previous buffer.
+  // The maximum number of residual bytes we require is 12; a five-byte
+  // unknown tag plus an eight-byte value, less one because the value
+  // is only a partial value.
+  char residual[12];
+  char *residual_end;
+
+  // Stores the user buffer passed to our decode function.
+  const char *buf_param;
+  size_t size_param;
+  const upb_bufhandle *handle;
+
+  // Our internal stack.
+  upb_pbdecoder_frame *stack, *top, *limit;
+  const uint32_t **callstack;
+  size_t stack_size;
+
+  upb_status *status;
+
+#ifdef UPB_USE_JIT_X64
+  // Used momentarily by the generated code to store a value while a user
+  // function is called.
+  uint32_t tmp_len;
+
+  const void *saved_rsp;
+#endif
+};
+
 // Decoder entry points; used as handlers.
 void *upb_pbdecoder_startbc(void *closure, const void *pc, size_t size_hint);
 void *upb_pbdecoder_startjit(void *closure, const void *hd, size_t size_hint);
diff --git a/upb/pb/encoder.c b/upb/pb/encoder.c
index d2c22e9..7a50996 100644
--- a/upb/pb/encoder.c
+++ b/upb/pb/encoder.c
@@ -62,6 +62,68 @@
 
 #include <stdlib.h>
 
+// The output buffer is divided into segments; a segment is a string of data
+// that is "ready to go" -- it does not need any varint lengths inserted into
+// the middle.  The seams between segments are where varints will be inserted
+// once they are known.
+//
+// We also use the concept of a "run", which is a range of encoded bytes that
+// occur at a single submessage level.  Every segment contains one or more runs.
+//
+// A segment can span messages.  Consider:
+//
+//                  .--Submessage lengths---------.
+//                  |       |                     |
+//                  |       V                     V
+//                  V      | |---------------    | |-----------------
+// Submessages:    | |-----------------------------------------------
+// Top-level msg: ------------------------------------------------------------
+//
+// Segments:          -----   -------------------   -----------------
+// Runs:              *----   *--------------*---   *----------------
+// (* marks the start)
+//
+// Note that the top-level menssage is not in any segment because it does not
+// have any length preceding it.
+//
+// A segment is only interrupted when another length needs to be inserted.  So
+// observe how the second segment spans both the inner submessage and part of
+// the next enclosing message.
+typedef struct {
+  uint32_t msglen;  // The length to varint-encode before this segment.
+  uint32_t seglen;  // Length of the segment.
+} upb_pb_encoder_segment;
+
+struct upb_pb_encoder {
+  upb_env *env;
+
+  // Our input and output.
+  upb_sink input_;
+  upb_bytessink *output_;
+
+  // The "subclosure" -- used as the inner closure as part of the bytessink
+  // protocol.
+  void *subc;
+
+  // The output buffer and limit, and our current write position.  "buf"
+  // initially points to "initbuf", but is dynamically allocated if we need to
+  // grow beyond the initial size.
+  char *buf, *ptr, *limit;
+
+  // The beginning of the current run, or undefined if we are at the top level.
+  char *runbegin;
+
+  // The list of segments we are accumulating.
+  upb_pb_encoder_segment *segbuf, *segptr, *seglimit;
+
+  // The stack of enclosing submessages.  Each entry in the stack points to the
+  // segment where this submessage's length is being accumulated.
+  int *stack, *top, *stacklimit;
+
+  // Depth of startmsg/endmsg calls.
+  int depth;
+};
+
 /* low-level buffering ********************************************************/
 
 // Low-level functions for interacting with the output buffer.
@@ -80,24 +142,22 @@ static upb_pb_encoder_segment *top(upb_pb_encoder *e) {
 // e->ptr.  Returns false if the bytes could not be allocated.
 static bool reserve(upb_pb_encoder *e, size_t bytes) {
   if ((e->limit - e->ptr) < bytes) {
+    // Grow buffer.
     size_t needed = bytes + (e->ptr - e->buf);
     size_t old_size = e->limit - e->buf;
+
     size_t new_size = old_size;
+
     while (new_size < needed) {
       new_size *= 2;
     }
 
-    char *realloc_from = (e->buf == e->initbuf) ? NULL : e->buf;
-    char *new_buf = realloc(realloc_from, new_size);
+    char *new_buf = upb_env_realloc(e->env, e->buf, old_size, new_size);
 
     if (new_buf == NULL) {
       return false;
     }
 
-    if (realloc_from == NULL) {
-      memcpy(new_buf, e->initbuf, old_size);
-    }
-
     e->ptr = new_buf + (e->ptr - e->buf);
     e->runbegin = new_buf + (e->runbegin - e->buf);
     e->limit = new_buf + new_size;
@@ -166,21 +226,17 @@ static bool start_delim(upb_pb_encoder *e) {
     }
 
     if (++e->segptr == e->seglimit) {
-      upb_pb_encoder_segment *realloc_from =
-          (e->segbuf == e->seginitbuf) ? NULL : e->segbuf;
+      // Grow segment buffer.
       size_t old_size =
           (e->seglimit - e->segbuf) * sizeof(upb_pb_encoder_segment);
       size_t new_size = old_size * 2;
-      upb_pb_encoder_segment *new_buf = realloc(realloc_from, new_size);
+      upb_pb_encoder_segment *new_buf =
+          upb_env_realloc(e->env, e->segbuf, old_size, new_size);
 
       if (new_buf == NULL) {
         return false;
       }
 
-      if (realloc_from == NULL) {
-        memcpy(new_buf, e->seginitbuf, old_size);
-      }
-
       e->segptr = new_buf + (e->segptr - e->segbuf);
       e->seglimit = new_buf + (new_size / sizeof(upb_pb_encoder_segment));
       e->segbuf = new_buf;
@@ -378,8 +434,10 @@ static void newhandlers_callback(const void *closure, upb_handlers *h) {
   upb_handlers_setendmsg(h, endmsg, NULL);
 
   const upb_msgdef *m = upb_handlers_msgdef(h);
-  upb_msg_iter i;
-  for(upb_msg_begin(&i, m); !upb_msg_done(&i); upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for(upb_msg_field_begin(&i, m);
+      !upb_msg_field_done(&i);
+      upb_msg_field_next(&i)) {
     const upb_fielddef *f = upb_msg_iter_field(&i);
     bool packed = upb_fielddef_isseq(f) && upb_fielddef_isprimitive(f) &&
                   upb_fielddef_packed(f);
@@ -449,6 +507,12 @@ static void newhandlers_callback(const void *closure, upb_handlers *h) {
   }
 }
 
+void upb_pb_encoder_reset(upb_pb_encoder *e) {
+  e->segptr = NULL;
+  e->top = NULL;
+  e->depth = 0;
+}
+
 
 /* public API *****************************************************************/
 
@@ -457,40 +521,42 @@ const upb_handlers *upb_pb_encoder_newhandlers(const upb_msgdef *m,
   return upb_handlers_newfrozen(m, owner, newhandlers_callback, NULL);
 }
 
-#define ARRAYSIZE(x) (sizeof(x) / sizeof(x[0]))
-
-void upb_pb_encoder_init(upb_pb_encoder *e, const upb_handlers *h) {
-  e->output_ = NULL;
-  e->subc = NULL;
-  e->buf = e->initbuf;
-  e->ptr = e->buf;
-  e->limit = e->buf + ARRAYSIZE(e->initbuf);
-  e->segbuf = e->seginitbuf;
-  e->seglimit = e->segbuf + ARRAYSIZE(e->seginitbuf);
-  e->stacklimit = e->stack + ARRAYSIZE(e->stack);
-  upb_sink_reset(&e->input_, h, e);
-}
-
-void upb_pb_encoder_uninit(upb_pb_encoder *e) {
-  if (e->buf != e->initbuf) {
-    free(e->buf);
+upb_pb_encoder *upb_pb_encoder_create(upb_env *env, const upb_handlers *h,
+                                      upb_bytessink *output) {
+  const size_t initial_bufsize = 256;
+  const size_t initial_segbufsize = 16;
+  // TODO(haberman): make this configurable.
+  const size_t stack_size = 64;
+#ifndef NDEBUG
+  const size_t size_before = upb_env_bytesallocated(env);
+#endif
+
+  upb_pb_encoder *e = upb_env_malloc(env, sizeof(upb_pb_encoder));
+  if (!e) return NULL;
+
+  e->buf = upb_env_malloc(env, initial_bufsize);
+  e->segbuf = upb_env_malloc(env, initial_segbufsize * sizeof(*e->segbuf));
+  e->stack = upb_env_malloc(env, stack_size * sizeof(*e->stack));
+
+  if (!e->buf || !e->segbuf || !e->stack) {
+    return NULL;
   }
 
-  if (e->segbuf != e->seginitbuf) {
-    free(e->segbuf);
-  }
-}
+  e->limit = e->buf + initial_bufsize;
+  e->seglimit = e->segbuf + initial_segbufsize;
+  e->stacklimit = e->stack + stack_size;
 
-void upb_pb_encoder_resetoutput(upb_pb_encoder *e, upb_bytessink *output) {
   upb_pb_encoder_reset(e);
+  upb_sink_reset(&e->input_, h, e);
+
+  e->env = env;
   e->output_ = output;
   e->subc = output->closure;
-}
+  e->ptr = e->buf;
 
-void upb_pb_encoder_reset(upb_pb_encoder *e) {
-  e->segptr = NULL;
-  e->top = NULL;
-  e->depth = 0;
+  // If this fails, increase the value in encoder.h.
+  assert(upb_env_bytesallocated(env) - size_before <= UPB_PB_ENCODER_SIZE);
+  return e;
 }
 
 upb_sink *upb_pb_encoder_input(upb_pb_encoder *e) { return &e->input_; }
diff --git a/upb/pb/encoder.h b/upb/pb/encoder.h
index 2df5797..edff95b 100644
--- a/upb/pb/encoder.h
+++ b/upb/pb/encoder.h
@@ -15,6 +15,7 @@
 #ifndef UPB_ENCODER_H_
 #define UPB_ENCODER_H_
 
+#include "upb/env.h"
 #include "upb/sink.h"
 
 #ifdef __cplusplus
@@ -31,101 +32,42 @@ UPB_DECLARE_TYPE(upb::pb::Encoder, upb_pb_encoder);
 
 /* upb::pb::Encoder ***********************************************************/
 
-// The output buffer is divided into segments; a segment is a string of data
-// that is "ready to go" -- it does not need any varint lengths inserted into
-// the middle.  The seams between segments are where varints will be inserted
-// once they are known.
-//
-// We also use the concept of a "run", which is a range of encoded bytes that
-// occur at a single submessage level.  Every segment contains one or more runs.
-//
-// A segment can span messages.  Consider:
-//
-//                  .--Submessage lengths---------.
-//                  |       |                     |
-//                  |       V                     V
-//                  V      | |---------------    | |-----------------
-// Submessages:    | |-----------------------------------------------
-// Top-level msg: ------------------------------------------------------------
-//
-// Segments:          -----   -------------------   -----------------
-// Runs:              *----   *--------------*---   *----------------
-// (* marks the start)
-//
-// Note that the top-level menssage is not in any segment because it does not
-// have any length preceding it.
-//
-// A segment is only interrupted when another length needs to be inserted.  So
-// observe how the second segment spans both the inner submessage and part of
-// the next enclosing message.
-typedef struct {
- UPB_PRIVATE_FOR_CPP
-  uint32_t msglen;  // The length to varint-encode before this segment.
-  uint32_t seglen;  // Length of the segment.
-} upb_pb_encoder_segment;
-
-UPB_DEFINE_CLASS0(upb::pb::Encoder,
- public:
-  Encoder(const upb::Handlers* handlers);
-  ~Encoder();
-
-  static reffed_ptr<const Handlers> NewHandlers(const upb::MessageDef* msg);
+// Preallocation hint: decoder won't allocate more bytes than this when first
+// constructed.  This hint may be an overestimate for some build configurations.
+// But if the decoder library is upgraded without recompiling the application,
+// it may be an underestimate.
+#define UPB_PB_ENCODER_SIZE 768
 
-  // Resets the state of the printer, so that it will expect to begin a new
-  // document.
-  void Reset();
+#ifdef __cplusplus
 
-  // Resets the output pointer which will serve as our closure.
-  void ResetOutput(BytesSink* output);
+class upb::pb::Encoder {
+ public:
+  // Creates a new encoder in the given environment.  The Handlers must have
+  // come from NewHandlers() below.
+  static Encoder* Create(Environment* env, const Handlers* handlers,
+                         BytesSink* output);
 
   // The input to the encoder.
   Sink* input();
 
- private:
-  UPB_DISALLOW_COPY_AND_ASSIGN(Encoder);
-,
-UPB_DEFINE_STRUCT0(upb_pb_encoder, UPB_QUOTE(
-  // Our input and output.
-  upb_sink input_;
-  upb_bytessink *output_;
-
-  // The "subclosure" -- used as the inner closure as part of the bytessink
-  // protocol.
-  void *subc;
-
-  // The output buffer and limit, and our current write position.  "buf"
-  // initially points to "initbuf", but is dynamically allocated if we need to
-  // grow beyond the initial size.
-  char *buf, *ptr, *limit;
+  // Creates a new set of handlers for this MessageDef.
+  static reffed_ptr<const Handlers> NewHandlers(const MessageDef* msg);
 
-  // The beginning of the current run, or undefined if we are at the top level.
-  char *runbegin;
+  static const size_t kSize = UPB_PB_ENCODER_SIZE;
 
-  // The list of segments we are accumulating.
-  upb_pb_encoder_segment *segbuf, *segptr, *seglimit;
-
-  // The stack of enclosing submessages.  Each entry in the stack points to the
-  // segment where this submessage's length is being accumulated.
-  int stack[UPB_PBENCODER_MAX_NESTING], *top, *stacklimit;
-
-  // Depth of startmsg/endmsg calls.
-  int depth;
+ private:
+  UPB_DISALLOW_POD_OPS(Encoder, upb::pb::Encoder);
+};
 
-  // Initial buffers for the output buffer and segment buffer.  If we outgrow
-  // these we will dynamically allocate bigger ones.
-  char initbuf[256];
-  upb_pb_encoder_segment seginitbuf[32];
-)));
+#endif
 
 UPB_BEGIN_EXTERN_C
 
 const upb_handlers *upb_pb_encoder_newhandlers(const upb_msgdef *m,
                                                const void *owner);
-void upb_pb_encoder_reset(upb_pb_encoder *e);
 upb_sink *upb_pb_encoder_input(upb_pb_encoder *p);
-void upb_pb_encoder_init(upb_pb_encoder *e, const upb_handlers *h);
-void upb_pb_encoder_resetoutput(upb_pb_encoder *e, upb_bytessink *output);
-void upb_pb_encoder_uninit(upb_pb_encoder *e);
+upb_pb_encoder* upb_pb_encoder_create(upb_env* e, const upb_handlers* h,
+                                      upb_bytessink* output);
 
 UPB_END_EXTERN_C
 
@@ -133,17 +75,9 @@ UPB_END_EXTERN_C
 
 namespace upb {
 namespace pb {
-inline Encoder::Encoder(const upb::Handlers* handlers) {
-  upb_pb_encoder_init(this, handlers);
-}
-inline Encoder::~Encoder() {
-  upb_pb_encoder_uninit(this);
-}
-inline void Encoder::Reset() {
-  upb_pb_encoder_reset(this);
-}
-inline void Encoder::ResetOutput(BytesSink* output) {
-  upb_pb_encoder_resetoutput(this, output);
+inline Encoder* Encoder::Create(Environment* env, const Handlers* handlers,
+                                BytesSink* output) {
+  return upb_pb_encoder_create(env, handlers, output);
 }
 inline Sink* Encoder::input() {
   return upb_pb_encoder_input(this);
diff --git a/upb/pb/glue.c b/upb/pb/glue.c
index fde2dd1..1259dac 100644
--- a/upb/pb/glue.c
+++ b/upb/pb/glue.c
@@ -22,26 +22,26 @@ upb_def **upb_load_defs_from_descriptor(const char *str, size_t len, int *n,
   const upb_pbdecodermethod *decoder_m =
       upb_pbdecodermethod_new(&opts, &decoder_m);
 
-  upb_pbdecoder decoder;
-  upb_descreader reader;
+  upb_env env;
+  upb_env_init(&env);
+  upb_env_reporterrorsto(&env, status);
 
-  upb_pbdecoder_init(&decoder, decoder_m, status);
-  upb_descreader_init(&reader, reader_h, status);
-  upb_pbdecoder_resetoutput(&decoder, upb_descreader_input(&reader));
+  upb_descreader *reader = upb_descreader_create(&env, reader_h);
+  upb_pbdecoder *decoder =
+      upb_pbdecoder_create(&env, decoder_m, upb_descreader_input(reader));
 
   // Push input data.
-  bool ok = upb_bufsrc_putbuf(str, len, upb_pbdecoder_input(&decoder));
+  bool ok = upb_bufsrc_putbuf(str, len, upb_pbdecoder_input(decoder));
 
   upb_def **ret = NULL;
 
   if (!ok) goto cleanup;
-  upb_def **defs = upb_descreader_getdefs(&reader, owner, n);
+  upb_def **defs = upb_descreader_getdefs(reader, owner, n);
   ret = malloc(sizeof(upb_def*) * (*n));
   memcpy(ret, defs, sizeof(upb_def*) * (*n));
 
 cleanup:
-  upb_pbdecoder_uninit(&decoder);
-  upb_descreader_uninit(&reader);
+  upb_env_uninit(&env);
   upb_handlers_unref(reader_h, &reader_h);
   upb_pbdecodermethod_unref(decoder_m, &decoder_m);
   return ret;
diff --git a/upb/pb/textprinter.c b/upb/pb/textprinter.c
index 8ceed68..07f951d 100644
--- a/upb/pb/textprinter.c
+++ b/upb/pb/textprinter.c
@@ -19,6 +19,14 @@
 
 #include "upb/sink.h"
 
+struct upb_textprinter {
+  upb_sink input_;
+  upb_bytessink *output_;
+  int indent_depth_;
+  bool single_line_;
+  void *subc;
+};
+
 #define CHECK(x) if ((x) < 0) goto err;
 
 static const char *shortname(const char *longname) {
@@ -236,24 +244,6 @@ err:
   return false;
 }
 
-
-/* Public API *****************************************************************/
-
-void upb_textprinter_init(upb_textprinter *p, const upb_handlers *h) {
-  p->single_line_ = false;
-  p->indent_depth_ = 0;
-  upb_sink_reset(&p->input_, h, p);
-}
-
-void upb_textprinter_uninit(upb_textprinter *p) {
-  UPB_UNUSED(p);
-}
-
-void upb_textprinter_reset(upb_textprinter *p, bool single_line) {
-  p->single_line_ = single_line;
-  p->indent_depth_ = 0;
-}
-
 static void onmreg(const void *c, upb_handlers *h) {
   UPB_UNUSED(c);
   const upb_msgdef *m = upb_handlers_msgdef(h);
@@ -261,8 +251,10 @@ static void onmreg(const void *c, upb_handlers *h) {
   upb_handlers_setstartmsg(h, textprinter_startmsg, NULL);
   upb_handlers_setendmsg(h, textprinter_endmsg, NULL);
 
-  upb_msg_iter i;
-  for(upb_msg_begin(&i, m); !upb_msg_done(&i); upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for(upb_msg_field_begin(&i, m);
+      !upb_msg_field_done(&i);
+      upb_msg_field_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
     upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
     upb_handlerattr_sethandlerdata(&attr, f);
@@ -311,6 +303,26 @@ static void onmreg(const void *c, upb_handlers *h) {
   }
 }
 
+static void textprinter_reset(upb_textprinter *p, bool single_line) {
+  p->single_line_ = single_line;
+  p->indent_depth_ = 0;
+}
+
+
+/* Public API *****************************************************************/
+
+upb_textprinter *upb_textprinter_create(upb_env *env, const upb_handlers *h,
+                                        upb_bytessink *output) {
+  upb_textprinter *p = upb_env_malloc(env, sizeof(upb_textprinter));
+  if (!p) return NULL;
+
+  p->output_ = output;
+  upb_sink_reset(&p->input_, h, p);
+  textprinter_reset(p, false);
+
+  return p;
+}
+
 const upb_handlers *upb_textprinter_newhandlers(const upb_msgdef *m,
                                                 const void *owner) {
   return upb_handlers_newfrozen(m, owner, &onmreg, NULL);
@@ -318,11 +330,6 @@ const upb_handlers *upb_textprinter_newhandlers(const upb_msgdef *m,
 
 upb_sink *upb_textprinter_input(upb_textprinter *p) { return &p->input_; }
 
-bool upb_textprinter_resetoutput(upb_textprinter *p, upb_bytessink *output) {
-  p->output_ = output;
-  return true;
-}
-
 void upb_textprinter_setsingleline(upb_textprinter *p, bool single_line) {
   p->single_line_ = single_line;
 }
diff --git a/upb/pb/textprinter.h b/upb/pb/textprinter.h
index 97e01f7..3ba3403 100644
--- a/upb/pb/textprinter.h
+++ b/upb/pb/textprinter.h
@@ -8,6 +8,7 @@
 #ifndef UPB_TEXT_H_
 #define UPB_TEXT_H_
 
+#include "upb/env.h"
 #include "upb/sink.h"
 
 #ifdef __cplusplus
@@ -20,58 +21,51 @@ class TextPrinter;
 
 UPB_DECLARE_TYPE(upb::pb::TextPrinter, upb_textprinter);
 
-UPB_DEFINE_CLASS0(upb::pb::TextPrinter,
+#ifdef __cplusplus
+
+class upb::pb::TextPrinter {
  public:
   // The given handlers must have come from NewHandlers().  It must outlive the
   // TextPrinter.
-  explicit TextPrinter(const upb::Handlers* handlers);
+  static TextPrinter *Create(Environment *env, const upb::Handlers *handlers,
+                             BytesSink *output);
 
   void SetSingleLineMode(bool single_line);
 
-  bool ResetOutput(BytesSink* output);
   Sink* input();
 
   // If handler caching becomes a requirement we can add a code cache as in
   // decoder.h
   static reffed_ptr<const Handlers> NewHandlers(const MessageDef* md);
+};
 
- private:
-,
-UPB_DEFINE_STRUCT0(upb_textprinter,
-  upb_sink input_;
-  upb_bytessink *output_;
-  int indent_depth_;
-  bool single_line_;
-  void *subc;
-));
+#endif
 
-UPB_BEGIN_EXTERN_C  // {
+UPB_BEGIN_EXTERN_C
 
 // C API.
-void upb_textprinter_init(upb_textprinter *p, const upb_handlers *h);
-void upb_textprinter_uninit(upb_textprinter *p);
-bool upb_textprinter_resetoutput(upb_textprinter *p, upb_bytessink *output);
+upb_textprinter *upb_textprinter_create(upb_env *env, const upb_handlers *h,
+                                        upb_bytessink *output);
 void upb_textprinter_setsingleline(upb_textprinter *p, bool single_line);
 upb_sink *upb_textprinter_input(upb_textprinter *p);
 
 const upb_handlers *upb_textprinter_newhandlers(const upb_msgdef *m,
                                                 const void *owner);
 
-UPB_END_EXTERN_C  // }
+UPB_END_EXTERN_C
 
 #ifdef __cplusplus
 
 namespace upb {
 namespace pb {
-inline TextPrinter::TextPrinter(const upb::Handlers* handlers) {
-  upb_textprinter_init(this, handlers);
+inline TextPrinter *TextPrinter::Create(Environment *env,
+                                        const upb::Handlers *handlers,
+                                        BytesSink *output) {
+  return upb_textprinter_create(env, handlers, output);
 }
 inline void TextPrinter::SetSingleLineMode(bool single_line) {
   upb_textprinter_setsingleline(this, single_line);
 }
-inline bool TextPrinter::ResetOutput(BytesSink* output) {
-  return upb_textprinter_resetoutput(this, output);
-}
 inline Sink* TextPrinter::input() {
   return upb_textprinter_input(this);
 }