1 files changed, 75 insertions, 85 deletions

diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index 5730504d..fe249d45 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -622,6 +622,49 @@ static const upb_pbdecodermethod *msgdef_decodermethod(Descriptor* desc) {
   return desc->fill_method;
 }
 
+
+// Stack-allocated context during an encode/decode operation. Contains the upb
+// environment and its stack-based allocator, an initial buffer for allocations
+// to avoid malloc() when possible, and a template for Ruby exception messages
+// if any error occurs.
+#define STACK_ENV_STACKBYTES 4096
+typedef struct {
+  upb_env env;
+  upb_seededalloc alloc;
+  const char* ruby_error_template;
+  char allocbuf[STACK_ENV_STACKBYTES];
+} stackenv;
+
+static void stackenv_init(stackenv* se, const char* errmsg);
+static void stackenv_uninit(stackenv* se);
+
+// Callback invoked by upb if any error occurs during parsing or serialization.
+static bool env_error_func(void* ud, const upb_status* status) {
+  stackenv* se = ud;
+  // Free the env -- rb_raise will longjmp up the stack past the encode/decode
+  // function so it would not otherwise have been freed.
+  stackenv_uninit(se);
+  rb_raise(rb_eRuntimeError, se->ruby_error_template,
+           upb_status_errmsg(status));
+  // Never reached: rb_raise() always longjmp()s up the stack, past all of our
+  // code, back to Ruby.
+  return false;
+}
+
+static void stackenv_init(stackenv* se, const char* errmsg) {
+  se->ruby_error_template = errmsg;
+  upb_env_init(&se->env);
+  upb_seededalloc_init(&se->alloc, &se->allocbuf, STACK_ENV_STACKBYTES);
+  upb_env_setallocfunc(
+      &se->env, upb_seededalloc_getallocfunc(&se->alloc), &se->alloc);
+  upb_env_seterrorfunc(&se->env, env_error_func, se);
+}
+
+static void stackenv_uninit(stackenv* se) {
+  upb_env_uninit(&se->env);
+  upb_seededalloc_uninit(&se->alloc);
+}
+
 /*
  * call-seq:
  *     MessageClass.decode(data) => message
@@ -631,7 +674,7 @@ static const upb_pbdecodermethod *msgdef_decodermethod(Descriptor* desc) {
  * and returns a message object with the corresponding field values.
  */
 VALUE Message_decode(VALUE klass, VALUE data) {
-  VALUE descriptor = rb_iv_get(klass, kDescriptorInstanceVar);
+  VALUE descriptor = rb_ivar_get(klass, descriptor_instancevar_interned);
   Descriptor* desc = ruby_to_Descriptor(descriptor);
   VALUE msgklass = Descriptor_msgclass(descriptor);
 
@@ -645,21 +688,17 @@ VALUE Message_decode(VALUE klass, VALUE data) {
 
   const upb_pbdecodermethod* method = msgdef_decodermethod(desc);
   const upb_handlers* h = upb_pbdecodermethod_desthandlers(method);
-  upb_pbdecoder decoder;
-  upb_sink sink;
-  upb_status status = UPB_STATUS_INIT;
+  stackenv se;
+  stackenv_init(&se, "Error occurred during parsing: %s");
 
-  upb_pbdecoder_init(&decoder, method, &status);
+  upb_sink sink;
   upb_sink_reset(&sink, h, msg);
-  upb_pbdecoder_resetoutput(&decoder, &sink);
+  upb_pbdecoder* decoder =
+      upb_pbdecoder_create(&se.env, method, &sink);
   upb_bufsrc_putbuf(RSTRING_PTR(data), RSTRING_LEN(data),
-                    upb_pbdecoder_input(&decoder));
+                    upb_pbdecoder_input(decoder));
 
-  upb_pbdecoder_uninit(&decoder);
-  if (!upb_ok(&status)) {
-    rb_raise(rb_eRuntimeError, "Error occurred during parsing: %s.",
-             upb_status_errmsg(&status));
-  }
+  stackenv_uninit(&se);
 
   return msg_rb;
 }
@@ -673,7 +712,7 @@ VALUE Message_decode(VALUE klass, VALUE data) {
  * and returns a message object with the corresponding field values.
  */
 VALUE Message_decode_json(VALUE klass, VALUE data) {
-  VALUE descriptor = rb_iv_get(klass, kDescriptorInstanceVar);
+  VALUE descriptor = rb_ivar_get(klass, descriptor_instancevar_interned);
   Descriptor* desc = ruby_to_Descriptor(descriptor);
   VALUE msgklass = Descriptor_msgclass(descriptor);
 
@@ -688,21 +727,16 @@ VALUE Message_decode_json(VALUE klass, VALUE data) {
   MessageHeader* msg;
   TypedData_Get_Struct(msg_rb, MessageHeader, &Message_type, msg);
 
-  upb_status status = UPB_STATUS_INIT;
-  upb_json_parser parser;
-  upb_json_parser_init(&parser, &status);
+  stackenv se;
+  stackenv_init(&se, "Error occurred during parsing: %s");
 
   upb_sink sink;
   upb_sink_reset(&sink, get_fill_handlers(desc), msg);
-  upb_json_parser_resetoutput(&parser, &sink);
+  upb_json_parser* parser = upb_json_parser_create(&se.env, &sink);
   upb_bufsrc_putbuf(RSTRING_PTR(data), RSTRING_LEN(data),
-                    upb_json_parser_input(&parser));
+                    upb_json_parser_input(parser));
 
-  upb_json_parser_uninit(&parser);
-  if (!upb_ok(&status)) {
-    rb_raise(rb_eRuntimeError, "Error occurred during parsing: %s.",
-             upb_status_errmsg(&status));
-  }
+  stackenv_uninit(&se);
 
   return msg_rb;
 }
@@ -813,7 +847,7 @@ static void putsubmsg(VALUE submsg, const upb_fielddef *f, upb_sink *sink,
   if (submsg == Qnil) return;
 
   upb_sink subsink;
-  VALUE descriptor = rb_iv_get(submsg, kDescriptorInstanceVar);
+  VALUE descriptor = rb_ivar_get(submsg, descriptor_instancevar_interned);
   Descriptor* subdesc = ruby_to_Descriptor(descriptor);
 
   upb_sink_startsubmsg(sink, getsel(f, UPB_HANDLER_STARTSUBMSG), &subsink);
@@ -935,7 +969,8 @@ static void putmap(VALUE map, const upb_fielddef *f, upb_sink *sink,
     VALUE value = Map_iter_value(&it);
 
     upb_sink entry_sink;
-    upb_sink_startsubmsg(&subsink, getsel(f, UPB_HANDLER_STARTSUBMSG), &entry_sink);
+    upb_sink_startsubmsg(&subsink, getsel(f, UPB_HANDLER_STARTSUBMSG),
+                         &entry_sink);
     upb_sink_startmsg(&entry_sink);
 
     put_ruby_value(key, key_field, Qnil, depth + 1, &entry_sink);
@@ -956,7 +991,7 @@ static void putmsg(VALUE msg_rb, const Descriptor* desc,
 
   // Protect against cycles (possible because users may freely reassign message
   // and repeated fields) by imposing a maximum recursion depth.
-  if (depth > UPB_SINK_MAX_NESTING) {
+  if (depth > ENCODE_MAX_NESTING) {
     rb_raise(rb_eRuntimeError,
              "Maximum recursion depth exceeded during encoding.");
   }
@@ -1065,7 +1100,7 @@ static const upb_handlers* msgdef_json_serialize_handlers(Descriptor* desc) {
  * wire format.
  */
 VALUE Message_encode(VALUE klass, VALUE msg_rb) {
-  VALUE descriptor = rb_iv_get(klass, kDescriptorInstanceVar);
+  VALUE descriptor = rb_ivar_get(klass, descriptor_instancevar_interned);
   Descriptor* desc = ruby_to_Descriptor(descriptor);
 
   stringsink sink;
@@ -1074,15 +1109,16 @@ VALUE Message_encode(VALUE klass, VALUE msg_rb) {
   const upb_handlers* serialize_handlers =
       msgdef_pb_serialize_handlers(desc);
 
-  upb_pb_encoder encoder;
-  upb_pb_encoder_init(&encoder, serialize_handlers);
-  upb_pb_encoder_resetoutput(&encoder, &sink.sink);
+  stackenv se;
+  stackenv_init(&se, "Error occurred during encoding: %s");
+  upb_pb_encoder* encoder =
+      upb_pb_encoder_create(&se.env, serialize_handlers, &sink.sink);
 
-  putmsg(msg_rb, desc, upb_pb_encoder_input(&encoder), 0);
+  putmsg(msg_rb, desc, upb_pb_encoder_input(encoder), 0);
 
   VALUE ret = rb_str_new(sink.ptr, sink.len);
 
-  upb_pb_encoder_uninit(&encoder);
+  stackenv_uninit(&se);
   stringsink_uninit(&sink);
 
   return ret;
@@ -1095,7 +1131,7 @@ VALUE Message_encode(VALUE klass, VALUE msg_rb) {
  * Encodes the given message object into its serialized JSON representation.
  */
 VALUE Message_encode_json(VALUE klass, VALUE msg_rb) {
-  VALUE descriptor = rb_iv_get(klass, kDescriptorInstanceVar);
+  VALUE descriptor = rb_ivar_get(klass, descriptor_instancevar_interned);
   Descriptor* desc = ruby_to_Descriptor(descriptor);
 
   stringsink sink;
@@ -1104,64 +1140,18 @@ VALUE Message_encode_json(VALUE klass, VALUE msg_rb) {
   const upb_handlers* serialize_handlers =
       msgdef_json_serialize_handlers(desc);
 
-  upb_json_printer printer;
-  upb_json_printer_init(&printer, serialize_handlers);
-  upb_json_printer_resetoutput(&printer, &sink.sink);
+  stackenv se;
+  stackenv_init(&se, "Error occurred during encoding: %s");
+  upb_json_printer* printer =
+      upb_json_printer_create(&se.env, serialize_handlers, &sink.sink);
 
-  putmsg(msg_rb, desc, upb_json_printer_input(&printer), 0);
+  putmsg(msg_rb, desc, upb_json_printer_input(printer), 0);
 
   VALUE ret = rb_str_new(sink.ptr, sink.len);
 
-  upb_json_printer_uninit(&printer);
+  stackenv_uninit(&se);
   stringsink_uninit(&sink);
 
   return ret;
 }
 
-/*
- * call-seq:
- *     Google::Protobuf.encode(msg) => bytes
- *
- * Encodes the given message object to protocol buffers wire format. This is an
- * alternative to the #encode method on msg's class.
- */
-VALUE Google_Protobuf_encode(VALUE self, VALUE msg_rb) {
-  VALUE klass = CLASS_OF(msg_rb);
-  return Message_encode(klass, msg_rb);
-}
-
-/*
- * call-seq:
- *     Google::Protobuf.encode_json(msg) => json_string
- *
- * Encodes the given message object to its JSON representation. This is an
- * alternative to the #encode_json method on msg's class.
- */
-VALUE Google_Protobuf_encode_json(VALUE self, VALUE msg_rb) {
-  VALUE klass = CLASS_OF(msg_rb);
-  return Message_encode_json(klass, msg_rb);
-}
-
-/*
- * call-seq:
- *     Google::Protobuf.decode(class, bytes) => msg
- *
- * Decodes the given bytes as protocol buffers wire format under the
- * interpretation given by the given class's message definition. This is an
- * alternative to the #decode method on the given class.
- */
-VALUE Google_Protobuf_decode(VALUE self, VALUE klass, VALUE msg_rb) {
-  return Message_decode(klass, msg_rb);
-}
-
-/*
- * call-seq:
- *     Google::Protobuf.decode_json(class, json_string) => msg
- *
- * Decodes the given JSON string under the interpretation given by the given
- * class's message definition. This is an alternative to the #decode_json method
- * on the given class.
- */
-VALUE Google_Protobuf_decode_json(VALUE self, VALUE klass, VALUE msg_rb) {
-  return Message_decode_json(klass, msg_rb);
-}