WIP.

cmake/google/protobuf/descriptor.upb.c

@@ -126,16 +126,18 @@ const upb_msglayout google_protobuf_ExtensionRangeOptions_msginit = {
   UPB_SIZE(8, 8), 1, _UPB_MSGEXT_EXTENDABLE, 0, 255,
 };
 
-static const upb_msglayout_sub google_protobuf_FieldDescriptorProto_submsgs[1] = {
+static const upb_msglayout_sub google_protobuf_FieldDescriptorProto_submsgs[3] = {
   {.submsg = &google_protobuf_FieldOptions_msginit},
+  {.subenum = &google_protobuf_FieldDescriptorProto_Label_enuminit},
+  {.subenum = &google_protobuf_FieldDescriptorProto_Type_enuminit},
 };
 
 static const upb_msglayout_field google_protobuf_FieldDescriptorProto__fields[11] = {
   {1, UPB_SIZE(24, 24), 1, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {2, UPB_SIZE(32, 40), 2, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {3, UPB_SIZE(12, 12), 3, 0, 5, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
-  {4, UPB_SIZE(4, 4), 4, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {4, UPB_SIZE(4, 4), 4, 1, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
-  {5, UPB_SIZE(8, 8), 5, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {5, UPB_SIZE(8, 8), 5, 2, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
   {6, UPB_SIZE(40, 56), 6, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {7, UPB_SIZE(48, 72), 7, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {8, UPB_SIZE(64, 104), 8, 0, 11, _UPB_MODE_SCALAR | (_UPB_REP_PTR << _UPB_REP_SHIFT)},
@@ -248,14 +250,15 @@ const upb_msglayout google_protobuf_MethodDescriptorProto_msginit = {
   UPB_SIZE(32, 64), 6, _UPB_MSGEXT_NONE, 6, 255,
 };
 
-static const upb_msglayout_sub google_protobuf_FileOptions_submsgs[1] = {
+static const upb_msglayout_sub google_protobuf_FileOptions_submsgs[2] = {
   {.submsg = &google_protobuf_UninterpretedOption_msginit},
+  {.subenum = &google_protobuf_FileOptions_OptimizeMode_enuminit},
 };
 
 static const upb_msglayout_field google_protobuf_FileOptions__fields[21] = {
   {1, UPB_SIZE(20, 24), 1, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {8, UPB_SIZE(28, 40), 2, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
-  {9, UPB_SIZE(4, 4), 3, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {9, UPB_SIZE(4, 4), 3, 1, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
   {10, UPB_SIZE(8, 8), 4, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
   {11, UPB_SIZE(36, 56), 5, 0, 12, _UPB_MODE_SCALAR | (_UPB_REP_STRVIEW << _UPB_REP_SHIFT)},
   {16, UPB_SIZE(9, 9), 6, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
@@ -300,16 +303,18 @@ const upb_msglayout google_protobuf_MessageOptions_msginit = {
   UPB_SIZE(16, 16), 5, _UPB_MSGEXT_EXTENDABLE, 3, 255,
 };
 
-static const upb_msglayout_sub google_protobuf_FieldOptions_submsgs[1] = {
+static const upb_msglayout_sub google_protobuf_FieldOptions_submsgs[3] = {
   {.submsg = &google_protobuf_UninterpretedOption_msginit},
+  {.subenum = &google_protobuf_FieldOptions_CType_enuminit},
+  {.subenum = &google_protobuf_FieldOptions_JSType_enuminit},
 };
 
 static const upb_msglayout_field google_protobuf_FieldOptions__fields[7] = {
-  {1, UPB_SIZE(4, 4), 1, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {1, UPB_SIZE(4, 4), 1, 1, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
   {2, UPB_SIZE(12, 12), 2, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
   {3, UPB_SIZE(13, 13), 3, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
   {5, UPB_SIZE(14, 14), 4, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
-  {6, UPB_SIZE(8, 8), 5, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {6, UPB_SIZE(8, 8), 5, 2, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
   {10, UPB_SIZE(15, 15), 6, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
   {999, UPB_SIZE(16, 16), 0, 0, 11, _UPB_MODE_ARRAY | (_UPB_REP_PTR << _UPB_REP_SHIFT)},
 };
@@ -380,13 +385,14 @@ const upb_msglayout google_protobuf_ServiceOptions_msginit = {
   UPB_SIZE(8, 16), 2, _UPB_MSGEXT_EXTENDABLE, 0, 255,
 };
 
-static const upb_msglayout_sub google_protobuf_MethodOptions_submsgs[1] = {
+static const upb_msglayout_sub google_protobuf_MethodOptions_submsgs[2] = {
   {.submsg = &google_protobuf_UninterpretedOption_msginit},
+  {.subenum = &google_protobuf_MethodOptions_IdempotencyLevel_enuminit},
 };
 
 static const upb_msglayout_field google_protobuf_MethodOptions__fields[3] = {
   {33, UPB_SIZE(8, 8), 1, 0, 8, _UPB_MODE_SCALAR | (_UPB_REP_1BYTE << _UPB_REP_SHIFT)},
-  {34, UPB_SIZE(4, 4), 2, 0, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
+  {34, UPB_SIZE(4, 4), 2, 1, 14, _UPB_MODE_SCALAR | (_UPB_REP_4BYTE << _UPB_REP_SHIFT)},
   {999, UPB_SIZE(12, 16), 0, 0, 11, _UPB_MODE_ARRAY | (_UPB_REP_PTR << _UPB_REP_SHIFT)},
 };
 
@@ -512,10 +518,57 @@ static const upb_msglayout *messages_layout[27] = {
   &google_protobuf_GeneratedCodeInfo_Annotation_msginit,
 };
 
+const upb_enumlayout google_protobuf_FieldDescriptorProto_Label_enuminit = {
+  NULL,
+  0xeULL,
+  0,
+};
+
+const upb_enumlayout google_protobuf_FieldDescriptorProto_Type_enuminit = {
+  NULL,
+  0x7fffeULL,
+  0,
+};
+
+const upb_enumlayout google_protobuf_FieldOptions_CType_enuminit = {
+  NULL,
+  0x7ULL,
+  0,
+};
+
+const upb_enumlayout google_protobuf_FieldOptions_JSType_enuminit = {
+  NULL,
+  0x7ULL,
+  0,
+};
+
+const upb_enumlayout google_protobuf_FileOptions_OptimizeMode_enuminit = {
+  NULL,
+  0xeULL,
+  0,
+};
+
+const upb_enumlayout google_protobuf_MethodOptions_IdempotencyLevel_enuminit = {
+  NULL,
+  0x7ULL,
+  0,
+};
+
+static const upb_enumlayout *enums_layout[6] = {
+  &google_protobuf_FieldDescriptorProto_Label_enuminit,
+  &google_protobuf_FieldDescriptorProto_Type_enuminit,
+  &google_protobuf_FieldOptions_CType_enuminit,
+  &google_protobuf_FieldOptions_JSType_enuminit,
+  &google_protobuf_FileOptions_OptimizeMode_enuminit,
+  &google_protobuf_MethodOptions_IdempotencyLevel_enuminit,
+};
+
 const upb_msglayout_file google_protobuf_descriptor_proto_upb_file_layout = {
   messages_layout,
+  enums_layout,
   NULL,
   27,
+  6,
   0,
 };
 

cmake/google/protobuf/descriptor.upb.h

@@ -154,6 +154,13 @@ typedef enum {
 } google_protobuf_MethodOptions_IdempotencyLevel;
 
 
+extern const upb_enumlayout google_protobuf_FieldDescriptorProto_Label_enuminit;
+extern const upb_enumlayout google_protobuf_FieldDescriptorProto_Type_enuminit;
+extern const upb_enumlayout google_protobuf_FieldOptions_CType_enuminit;
+extern const upb_enumlayout google_protobuf_FieldOptions_JSType_enuminit;
+extern const upb_enumlayout google_protobuf_FileOptions_OptimizeMode_enuminit;
+extern const upb_enumlayout google_protobuf_MethodOptions_IdempotencyLevel_enuminit;
+
 /* google.protobuf.FileDescriptorSet */
 
 UPB_INLINE google_protobuf_FileDescriptorSet *google_protobuf_FileDescriptorSet_new(upb_arena *arena) {

upb/msg_internal.h

@@ -66,8 +66,8 @@ enum {
 typedef struct {
   uint32_t number;
   uint16_t offset;
-  int16_t presence;       /* If >0, hasbit_index.  If <0, ~oneof_index. */
+  int16_t presence;       // If >0, hasbit_index.  If <0, ~oneof_index
-  uint16_t submsg_index;  /* undefined if descriptortype != MESSAGE or GROUP. */
+  uint16_t submsg_index;  // undefined if descriptortype != MESSAGE/GROUP/ENUM
   uint8_t descriptortype;
   uint8_t mode; /* upb_fieldmode | upb_labelflags |
                    (upb_rep << _UPB_REP_SHIFT) */
@@ -130,9 +130,15 @@ typedef struct {
   _upb_field_parser *field_parser;
 } _upb_fasttable_entry;
 
+typedef struct {
+  const int32_t *values;  // List of values <0 or >63
+  uint64_t mask;          // Bits are set for acceptable value 0 <= x < 64
+  int value_count;
+} upb_enumlayout;
+
 typedef union {
   const struct upb_msglayout *submsg;
-  // TODO: const upb_enumlayout *subenum;
+  const upb_enumlayout *subenum;
 } upb_msglayout_sub;
 
 typedef enum {
@@ -179,8 +185,10 @@ typedef struct {
 
 typedef struct {
   const upb_msglayout **msgs;
+  const upb_enumlayout **enums;
   const upb_msglayout_ext **exts;
   int msg_count;
+  int enum_count;
   int ext_count;
 } upb_msglayout_file;
 

upb/msg_test.cc

@@ -26,6 +26,7 @@
  */
 
 #include "gtest/gtest.h"
+#include "gmock/gmock.h"
 #include "src/google/protobuf/test_messages_proto3.upb.h"
 #include "upb/def.hpp"
 #include "upb/json_decode.h"
@@ -152,3 +153,52 @@ TEST(MessageTest, MessageSet) {
       << status.error_message();
   VerifyMessageSet(ext_msg3);
 }
+
+TEST(MessageTest, Proto2Enum) {
+  upb::Arena arena;
+  upb_test_Proto2FakeEnumMessage *fake_msg = upb_test_Proto2FakeEnumMessage_new(arena.ptr());
+
+  upb_test_Proto2FakeEnumMessage_set_optional_enum(fake_msg, 999);
+  int32_t *vals = upb_test_Proto2FakeEnumMessage_resize_repeated_enum(
+      fake_msg, 6, arena.ptr());
+  vals[0] = upb_test_Proto2EnumMessage_ZERO;
+  vals[1] = 7;  // Unknown small.
+  vals[2] = upb_test_Proto2EnumMessage_SMALL;
+  vals[3] = 888;  // Unknown large.
+  vals[4] = upb_test_Proto2EnumMessage_LARGE;
+  vals[5] = upb_test_Proto2EnumMessage_NEGATIVE;
+  upb_test_Proto2FakeEnumMessage_int32_enum_map_set(fake_msg, 5, 777,
+                                                    arena.ptr());
+
+  size_t size;
+  char *pb =
+      upb_test_Proto2FakeEnumMessage_serialize(fake_msg, arena.ptr(), &size);
+
+  upb_test_Proto2EnumMessage *enum_msg =
+      upb_test_Proto2EnumMessage_parse(pb, size, arena.ptr());
+  ASSERT_TRUE(enum_msg != nullptr);
+
+  EXPECT_EQ(false, upb_test_Proto2EnumMessage_has_optional_enum(enum_msg));
+  const int32_t *vals_const =
+      upb_test_Proto2EnumMessage_repeated_enum(enum_msg, &size);
+  EXPECT_EQ(4, size);  // Two unknown values moved to the unknown field set.
+
+  pb = upb_test_Proto2EnumMessage_serialize(enum_msg, arena.ptr(), &size);
+  upb_test_Proto2FakeEnumMessage *fake_msg2 =
+      upb_test_Proto2FakeEnumMessage_parse(pb, size, arena.ptr());
+
+  EXPECT_EQ(true, upb_test_Proto2FakeEnumMessage_has_optional_enum(fake_msg2));
+  vals_const =
+      upb_test_Proto2FakeEnumMessage_repeated_enum(fake_msg2, &size);
+  EXPECT_EQ(6, size);
+  int32_t expected[] = {
+      upb_test_Proto2EnumMessage_ZERO,
+      upb_test_Proto2EnumMessage_SMALL,
+      upb_test_Proto2EnumMessage_LARGE,
+      upb_test_Proto2EnumMessage_NEGATIVE,
+      7,
+      888,
+  };
+  EXPECT_THAT(std::vector<int32_t>(vals_const, vals_const + size),
+              ::testing::ElementsAreArray(expected));
+}

upb/msg_test.proto

@@ -58,3 +58,24 @@ message MessageSetMember {
     optional MessageSetMember message_set_extension = 4;
   }
 }
+
+message Proto2EnumMessage {
+  enum Proto2TestEnum {
+    ZERO = 0;
+    NEGATIVE = -1;
+    SMALL = 15;
+    LARGE = 12345;
+  }
+
+  optional Proto2TestEnum optional_enum = 1;
+  repeated Proto2TestEnum repeated_enum = 2;
+  map<int32, Proto2TestEnum> int32_enum_map = 3;
+}
+
+// The same fields as Proto2EnumMessage, but with int32 fields so we can fake
+// wire format.
+message Proto2FakeEnumMessage {
+  optional int32 optional_enum = 1;
+  repeated int32 repeated_enum = 2;
+  map<int32, int32> int32_enum_map = 3;
+}

upbc/protoc-gen-upb.cc

@@ -51,6 +51,10 @@ std::string MessageInit(const protobuf::Descriptor* descriptor) {
   return MessageName(descriptor) + "_msginit";
 }
 
+std::string EnumInit(const protobuf::EnumDescriptor* descriptor) {
+  return ToCIdent(descriptor->full_name()) + "_enuminit";
+}
+
 std::string ExtensionIdentBase(const protobuf::FieldDescriptor* ext) {
   assert(ext->is_extension());
   std::string ext_scope;
@@ -65,8 +69,9 @@ std::string ExtensionLayout(const google::protobuf::FieldDescriptor* ext) {
   return absl::StrCat(ExtensionIdentBase(ext), "_", ext->name(), "_ext");
 }
 
-const char *kMessagesInit = "messages_layout";
+const char *kEnumsInit = "enums_layout";
 const char *kExtensionsInit = "extensions_layout";
+const char *kMessagesInit = "messages_layout";
 
 void AddEnums(const protobuf::Descriptor* message,
               std::vector<const protobuf::EnumDescriptor*>* enums) {
@@ -179,6 +184,24 @@ std::vector<const protobuf::FieldDescriptor*> SortedSubmessages(
   return ret;
 }
 
+std::vector<const protobuf::FieldDescriptor*> SortedSubEnums(
+    const protobuf::Descriptor* message) {
+  std::vector<const protobuf::FieldDescriptor*> ret;
+  for (int i = 0; i < message->field_count(); i++) {
+    if (message->field(i)->cpp_type() ==
+        protobuf::FieldDescriptor::CPPTYPE_ENUM) {
+      ret.push_back(message->field(i));
+    }
+  }
+  std::sort(ret.begin(), ret.end(),
+            [](const protobuf::FieldDescriptor* a,
+               const protobuf::FieldDescriptor* b) {
+              return a->enum_type()->full_name() <
+                     b->enum_type()->full_name();
+            });
+  return ret;
+}
+
 std::string EnumValueSymbol(const protobuf::EnumValueDescriptor* value) {
   return ToCIdent(value->full_name());
 }
@@ -721,6 +744,14 @@ void WriteHeader(const protobuf::FileDescriptor* file, Output& output) {
 
   output("\n");
 
+  if (file->syntax() == protobuf::FileDescriptor::SYNTAX_PROTO2) {
+    for (auto enumdesc : this_file_enums) {
+      output("extern const upb_enumlayout $0;\n", EnumInit(enumdesc));
+    }
+  }
+
+  output("\n");
+
   for (auto message : this_file_messages) {
     GenerateMessageInHeader(message, output);
   }
@@ -787,45 +818,70 @@ int TableDescriptorType(const protobuf::FieldDescriptor* field) {
     // embedding field names on the side) we will have to revisit this, because
     // string vs. bytes behavior is not affected by proto2 vs proto3.
     return protobuf::FieldDescriptor::TYPE_BYTES;
+  } else if (field->enum_type() &&
+             field->enum_type()->file()->syntax() ==
+                 protobuf::FileDescriptor::SYNTAX_PROTO3) {
+    // From the perspective of the binary decoder, proto3 enums are identical to
+    // int32 fields.  Only in proto2 do we check enum values to make sure they
+    // are defined in the enum.
+    return protobuf::FieldDescriptor::TYPE_INT32;
   } else {
     return field->type();
   }
 }
 
-struct SubmsgArray {
+struct SubLayoutArray {
  public:
-  SubmsgArray(const protobuf::Descriptor* message) : message_(message) {
+  SubLayoutArray(const protobuf::Descriptor* message);
-    MessageLayout layout(message);
-    std::vector<const protobuf::FieldDescriptor*> sorted_submsgs =
-        SortedSubmessages(message);
-    int i = 0;
-    for (auto submsg : sorted_submsgs) {
-      if (indexes_.find(submsg->message_type()) != indexes_.end()) {
-        continue;
-      }
-      submsgs_.push_back(submsg->message_type());
-      indexes_[submsg->message_type()] = i++;
-    }
-  }
 
   const std::vector<const protobuf::Descriptor*>& submsgs() const {
     return submsgs_;
   }
 
-  int GetIndex(const protobuf::FieldDescriptor* field) {
+  const std::vector<const protobuf::EnumDescriptor*>& subenums() const {
-    (void)message_;
+    return subenums_;
-    assert(field->containing_type() == message_);
+  }
-    auto it = indexes_.find(field->message_type());
+
+  int total_count() const { return submsgs_.size() + subenums_.size(); }
+
+  int GetIndex(const void *sub) {
+    auto it = indexes_.find(sub);
     assert(it != indexes_.end());
     return it->second;
   }
 
  private:
-  const protobuf::Descriptor* message_;
   std::vector<const protobuf::Descriptor*> submsgs_;
-  absl::flat_hash_map<const protobuf::Descriptor*, int> indexes_;
+  std::vector<const protobuf::EnumDescriptor*> subenums_;
+  absl::flat_hash_map<const void*, int> indexes_;
 };
 
+SubLayoutArray::SubLayoutArray(const protobuf::Descriptor* message) {
+  MessageLayout layout(message);
+  std::vector<const protobuf::FieldDescriptor*> sorted_submsgs =
+      SortedSubmessages(message);
+  int i = 0;
+  for (auto submsg : sorted_submsgs) {
+    if (indexes_.find(submsg->message_type()) != indexes_.end()) {
+      continue;
+    }
+    submsgs_.push_back(submsg->message_type());
+    indexes_[submsg->message_type()] = i++;
+  }
+
+  std::vector<const protobuf::FieldDescriptor*> sorted_subenums =
+      SortedSubEnums(message);
+  for (auto field : sorted_subenums) {
+    if (field->file()->syntax() !=
+            protobuf::FileDescriptor::SYNTAX_PROTO2 ||
+        indexes_.find(field->enum_type()) != indexes_.end()) {
+      continue;
+    }
+    subenums_.push_back(field->enum_type());
+    indexes_[field->enum_type()] = i++;
+  }
+}
+
 typedef std::pair<std::string, uint64_t> TableEntry;
 
 uint64_t GetEncodedTag(const protobuf::FieldDescriptor* field) {
@@ -957,8 +1013,8 @@ bool TryFillTableEntry(const protobuf::Descriptor* message,
   }
 
   if (field->cpp_type() == protobuf::FieldDescriptor::CPPTYPE_MESSAGE) {
-    SubmsgArray submsg_array(message);
+    SubLayoutArray sublayout_array(message);
-    uint64_t idx = submsg_array.GetIndex(field);
+    uint64_t idx = sublayout_array.GetIndex(field->message_type());
     if (idx > 255) return false;
     data |= idx << 16;
 
@@ -1083,19 +1139,22 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output,
   uint8_t dense_below = 0;
   const int dense_below_max = std::numeric_limits<decltype(dense_below)>::max();
   MessageLayout layout(message);
-  SubmsgArray submsg_array(message);
+  SubLayoutArray sublayout_array(message);
 
-  if (!submsg_array.submsgs().empty()) {
+  if (sublayout_array.total_count()) {
     // TODO(haberman): could save a little bit of space by only generating a
     // "submsgs" array for every strongly-connected component.
     std::string submsgs_array_name = msg_name + "_submsgs";
     submsgs_array_ref = "&" + submsgs_array_name + "[0]";
     output("static const upb_msglayout_sub $0[$1] = {\n",
-           submsgs_array_name, submsg_array.submsgs().size());
+           submsgs_array_name, sublayout_array.total_count());
 
-    for (auto submsg : submsg_array.submsgs()) {
+    for (const auto* submsg : sublayout_array.submsgs()) {
       output("  {.submsg = &$0},\n", MessageInit(submsg));
     }
+    for (const auto* subenum : sublayout_array.subenums()) {
+      output("  {.subenum = &$0},\n", EnumInit(subenum));
+    }
 
     output("};\n\n");
   }
@@ -1109,7 +1168,7 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output,
            fields_array_name, field_number_order.size());
     for (int i = 0; i < static_cast<int>(field_number_order.size()); i++) {
       auto field = field_number_order[i];
-      int submsg_index = 0;
+      int sublayout_index = 0;
 
       if (i < dense_below_max && field->number() == i + 1 &&
           (i == 0 || field_number_order[i - 1]->number() == i)) {
@@ -1117,10 +1176,14 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output,
       }
 
       if (field->cpp_type() == protobuf::FieldDescriptor::CPPTYPE_MESSAGE) {
-        submsg_index = submsg_array.GetIndex(field);
+        sublayout_index = sublayout_array.GetIndex(field->message_type());
+      } else if (field->enum_type() &&
+                 field->enum_type()->file()->syntax() ==
+                     protobuf::FileDescriptor::SYNTAX_PROTO2) {
+        sublayout_index = sublayout_array.GetIndex(field->enum_type());
       }
 
-      WriteMessageField(field, layout, submsg_index, output);
+      WriteMessageField(field, layout, sublayout_index, output);
     }
     output("};\n\n");
   }
@@ -1168,6 +1231,61 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output,
   output("};\n\n");
 }
 
+int WriteEnums(const protobuf::FileDescriptor* file, Output& output) {
+  if (file->syntax() != protobuf::FileDescriptor::SYNTAX_PROTO2) {
+    return 0;
+  }
+
+  std::vector<const protobuf::EnumDescriptor*> this_file_enums =
+      SortedEnums(file);
+
+  std::string values_init = "NULL";
+
+  for (auto e : this_file_enums) {
+    uint64_t mask = 0;
+    absl::flat_hash_set<int32_t> values;
+    for (int i = 0; i < e->value_count(); i++) {
+      int32_t number = e->value(i)->number();
+      if (static_cast<uint32_t>(number) < 64) {
+        mask |= 1 << number;
+      } else {
+        values.insert(number);
+      }
+    }
+    std::vector<int32_t> values_vec(values.begin(), values.end());
+    std::sort(values_vec.begin(), values_vec.end());
+
+    if (!values_vec.empty()) {
+      values_init = EnumInit(e) + "_values";
+      output("static const int32_t $0[$1] = {\n", values_init,
+             values_vec.size());
+      for (auto value : values_vec) {
+        output("  $0,\n", value);
+      }
+      output("};\n\n");
+    }
+
+    output("const upb_enumlayout $0 = {\n", EnumInit(e));
+    output("  $0,\n", values_init);
+    output("  0x$0ULL,\n", absl::Hex(mask));
+    output("  $0,\n", values_vec.size());
+
+    output("};\n\n");
+  }
+
+  if (!this_file_enums.empty()) {
+    output("static const upb_enumlayout *$0[$1] = {\n", kEnumsInit,
+          this_file_enums.size());
+    for (auto e : this_file_enums) {
+      output("  &$0,\n", EnumInit(e));
+    }
+    output("};\n");
+    output("\n");
+  }
+
+  return this_file_enums.size();
+}
+
 void WriteExtension(const protobuf::FieldDescriptor* ext, Output& output) {
   output("const upb_msglayout_ext $0 = {\n  ", ExtensionLayout(ext));
   WriteField(ext, "0", "0", 0, output);
@@ -1264,11 +1382,14 @@ void WriteSource(const protobuf::FileDescriptor* file, Output& output,
 
   int msg_count = WriteMessages(file, output, fasttable_enabled);
   int ext_count = WriteExtensions(file, output);
+  int enum_count = WriteEnums(file, output);
 
   output("const upb_msglayout_file $0 = {\n", FileLayoutName(file));
   output("  $0,\n", msg_count ? kMessagesInit : "NULL");
+  output("  $0,\n", enum_count ? kEnumsInit : "NULL");
   output("  $0,\n", ext_count ? kExtensionsInit : "NULL");
   output("  $0,\n", msg_count);
+  output("  $0,\n", enum_count);
   output("  $0,\n", ext_count);
   output("};\n\n");