2 files changed, 37 insertions, 22 deletions

diff --git a/upb/bindings/lua/upb.c b/upb/bindings/lua/upb.c
index 17fc0a8..5ad0235 100644
--- a/upb/bindings/lua/upb.c
+++ b/upb/bindings/lua/upb.c
@@ -139,7 +139,8 @@ bool lupb_openlib(lua_State *L, void *ptr, const char *name,
 
 // Pushes a new userdata with the given metatable and ensures that it has a
 // uservalue.
-static void *newudata_with_userval(lua_State *L, size_t size, const char *type) {
+static void *newudata_with_userval(lua_State *L, size_t size,
+                                   const char *type) {
   void *ret = lua_newuserdata(L, size);
 
   // Set metatable.
@@ -952,17 +953,17 @@ static int lupb_msgdef_field(lua_State *L) {
 }
 
 static int lupb_msgiter_next(lua_State *L) {
-  upb_msg_iter *i = lua_touserdata(L, lua_upvalueindex(1));
-  if (upb_msg_done(i)) return 0;
+  upb_msg_field_iter *i = lua_touserdata(L, lua_upvalueindex(1));
+  if (upb_msg_field_done(i)) return 0;
   lupb_def_pushwrapper(L, UPB_UPCAST(upb_msg_iter_field(i)), NULL);
-  upb_msg_next(i);
+  upb_msg_field_next(i);
   return 1;
 }
 
 static int lupb_msgdef_fields(lua_State *L) {
   const upb_msgdef *m = lupb_msgdef_check(L, 1);
-  upb_msg_iter *i = lua_newuserdata(L, sizeof(upb_msg_iter));
-  upb_msg_begin(i, m);
+  upb_msg_field_iter *i = lua_newuserdata(L, sizeof(upb_msg_field_iter));
+  upb_msg_field_begin(i, m);
   // Need to guarantee that the msgdef outlives the iter.
   lua_pushvalue(L, 1);
   lua_pushcclosure(L, &lupb_msgiter_next, 2);
@@ -1416,8 +1417,10 @@ static lupb_msgdef *lupb_msg_assignoffsets(lua_State *L, int narg) {
   size_t userval_idx = 1;
 
   // Assign offsets.
-  upb_msg_iter i;
-  for (upb_msg_begin(&i, lmd->md); !upb_msg_done(&i); upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for (upb_msg_field_begin(&i, lmd->md);
+       !upb_msg_field_done(&i);
+       upb_msg_field_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
     if (in_userval(f)) {
       offsets[upb_fielddef_index(f)] = userval_idx++;
@@ -1442,7 +1445,9 @@ static lupb_msgdef *lupb_msg_assignoffsets(lua_State *L, int narg) {
   lua_newtable(L);  // This will be our userval.
 
   int idx = 1;
-  for (upb_msg_begin(&i, lmd->md); !upb_msg_done(&i); upb_msg_next(&i)) {
+  for (upb_msg_field_begin(&i, lmd->md);
+       !upb_msg_field_done(&i);
+       upb_msg_field_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
     if (upb_fielddef_type(f) == UPB_TYPE_MESSAGE) {
       bool created = lupb_def_pushwrapper(L, upb_fielddef_subdef(f), NULL);
@@ -1660,9 +1665,9 @@ void callback(const void *closure, upb_handlers *h) {
   lua_State *L = (lua_State*)closure;
   lupb_def_pushwrapper(L, UPB_UPCAST(upb_handlers_msgdef(h)), NULL);
   lupb_msgdef *lmd = lupb_msg_assignoffsets(L, -1);
-  upb_msg_iter i;
-  upb_msg_begin(&i, upb_handlers_msgdef(h));
-  for (; !upb_msg_done(&i); upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  upb_msg_field_begin(&i, upb_handlers_msgdef(h));
+  for (; !upb_msg_field_done(&i); upb_msg_field_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
     int hasbit = upb_fielddef_index(f);
     uint16_t ofs = lmd->field_offsets[upb_fielddef_index(f)];
diff --git a/upb/bindings/lua/upb/pb.c b/upb/bindings/lua/upb/pb.c
index c9f1f47..ea3d755 100644
--- a/upb/bindings/lua/upb/pb.c
+++ b/upb/bindings/lua/upb/pb.c
@@ -41,6 +41,13 @@ static int lupb_pbdecodermethod_new(lua_State *L) {
   return 1;  // The DecoderMethod wrapper.
 }
 
+// We implement upb's allocation function by allocating a Lua userdata.
+// This is a raw hunk of memory that will be GC'd by Lua.
+static void *lua_alloc(void *ud, size_t size) {
+  lua_State *L = ud;
+  return lua_newuserdata(L, size);
+}
+
 // Unlike most of our exposed Lua functions, this does not correspond to an
 // actual method on the underlying DecoderMethod.  But it's convenient, and
 // important to implement in C because we can do stack allocation and
@@ -61,19 +68,22 @@ static int lupb_pbdecodermethod_parse(lua_State *L) {
   // Handlers need this.
   lua_getuservalue(L, -1);
 
-  upb_pbdecoder decoder;
   upb_status status = UPB_STATUS_INIT;
-  upb_pbdecoder_init(&decoder, method, &status);
+  upb_env env;
+  upb_env_init(&env);
+  upb_env_reporterrorsto(&env, &status);
   upb_sink sink;
   upb_sink_reset(&sink, handlers, msg);
-  upb_pbdecoder_resetoutput(&decoder, &sink);
-  upb_bufsrc_putbuf(pb, len, upb_pbdecoder_input(&decoder));
-  // TODO: Our need to call uninit isn't longjmp-safe; what if the decode
-  // triggers a Lua error?  uninit is only needed if the decoder
-  // dynamically-allocated a growing stack -- ditch this feature and live with
-  // the compile-time limit?  Or have a custom allocation function that
-  // allocates Lua GC-rooted memory?
-  upb_pbdecoder_uninit(&decoder);
+  upb_pbdecoder *decoder = upb_pbdecoder_create(&env, method, &sink);
+  upb_bufsrc_putbuf(pb, len, upb_pbdecoder_input(decoder));
+
+  // This won't get called in the error case, which longjmp's across us.  But
+  // since we made our alloc function allocate only GC-able memory, that
+  // shouldn't matter.  It *would* matter if the environment had references to
+  // any non-memory resources (ie.  filehandles).  As an alternative to this we
+  // could make the environment itself a userdata.
+  upb_env_uninit(&env);
+
   lupb_checkstatus(L, &status);
 
   lua_pop(L, 1);  // Uservalue.