3 files changed, 19 insertions, 13 deletions
diff --git a/upb/bindings/googlepb/proto2.cc b/upb/bindings/googlepb/proto2.cc
index 3911172..87c13b6 100644
--- a/upb/bindings/googlepb/proto2.cc
+++ b/upb/bindings/googlepb/proto2.cc
@@ -946,14 +946,14 @@ case goog::FieldDescriptor::cpptype:                                           \
    public:
     typedef goog::Message Type;
 #ifdef GOOGLE_PROTOBUF_HAS_ARENAS
-    static ::proto2::Arena* GetArena(Type* t) {
+    static goog::Arena* GetArena(Type* t) {
       return t->GetArena();
     }
     static void* GetMaybeArenaPointer(Type* t) {
       return t->GetMaybeArenaPointer();
     }
     static inline Type* NewFromPrototype(
-        const Type* prototype, ::proto2::Arena* arena = NULL) {
+        const Type* prototype, goog::Arena* arena = NULL) {
       return prototype->New(arena);
     }
     static void Delete(Type* t, goog::Arena* arena = NULL) {
@@ -1277,6 +1277,11 @@ case goog::FieldDescriptor::cpptype:                                           \
       return lazy_field_.SetAllocated(static_cast<proto2::Message*>(message));
     }
 
+    virtual void UnsafeArenaSetAllocatedMessage(proto2::MessageLite* message) {
+      return lazy_field_.UnsafeArenaSetAllocated(
+          static_cast<proto2::Message*>(message));
+    }
+
     virtual proto2::MessageLite* ReleaseMessage(
         const proto2::MessageLite& prototype) {
       return lazy_field_.ReleaseByPrototype(
diff --git a/upb/bindings/lua/upb.c b/upb/bindings/lua/upb.c
index 5ad0235..b35af24 100644
--- a/upb/bindings/lua/upb.c
+++ b/upb/bindings/lua/upb.c
@@ -1358,7 +1358,7 @@ static size_t align_up(size_t val, size_t align) {
 
 // If we always read/write as a consistent type to each value, this shouldn't
 // violate aliasing.
-#define DEREF(msg, ofs, type) *(type*)(&msg->data[ofs])
+#define DEREF(msg, ofs, type) *(type*)((char*)msg + sizeof(lupb_msg) + ofs)
 
 lupb_msg *lupb_msg_check(lua_State *L, int narg) {
   lupb_msg *msg = luaL_checkudata(L, narg, LUPB_MSG);
diff --git a/upb/bindings/lua/upb/pb.c b/upb/bindings/lua/upb/pb.c
index c9f1f47..920648f 100644
--- a/upb/bindings/lua/upb/pb.c
+++ b/upb/bindings/lua/upb/pb.c
@@ -61,19 +61,20 @@ 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));
+
+  // TODO: This won't get called in the error case, which longjmp's across us.
+  // This will cause the memory to leak.  To remedy this, we should make the
+  // upb_env wrapped in a userdata that guarantees this will get called.
+  upb_env_uninit(&env);
+
   lupb_checkstatus(L, &status);
 
   lua_pop(L, 1);  // Uservalue.