From 74bc67bbbc77b059fcf65dfaa044d5e82035c781 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Mon, 21 Feb 2022 21:16:17 -0800
Subject: [PATCH] WIP

---
 upb/upb.h              |  2 +-
 upbc/message_layout.cc |  5 ++--
 upbc/message_layout.h  |  3 +++
 upbc/protoc-gen-upb.cc | 58 +++++++++++++++++++++++++++++++++++++-----
 4 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/upb/upb.h b/upb/upb.h
index 969e757eee..45000ffdaf 100644
--- a/upb/upb.h
+++ b/upb/upb.h
@@ -230,7 +230,7 @@ UPB_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize,
   oldsize = UPB_ALIGN_MALLOC(size);
   size = UPB_ALIGN_MALLOC(size);
   if (ch_ptr + oldsize == h->ptr) {
-    if (h->end - ch_ptr >= size) {
+    if ((size_t)(h->end - ch_ptr) >= size) {
       h->ptr = ch_ptr + size;
       return ptr;
     }
diff --git a/upbc/message_layout.cc b/upbc/message_layout.cc
index be63a105c8..e72dd6cfb2 100644
--- a/upbc/message_layout.cc
+++ b/upbc/message_layout.cc
@@ -43,8 +43,9 @@ MessageLayout::Size MessageLayout::Place(
   offset.AlignUp(size_and_align.align);
   size_ = offset;
   size_.Add(size_and_align.size);
-  // maxalign_.MaxFrom(size_and_align.align);
-  maxalign_.MaxFrom(size_and_align.size);
+  ABSL_ASSERT(IsPowerOfTwo(size_and_align.align.size32));
+  ABSL_ASSERT(IsPowerOfTwo(size_and_align.align.size64));
+  maxalign_.MaxFrom(size_and_align.align);
   return offset;
 }
 
diff --git a/upbc/message_layout.h b/upbc/message_layout.h
index cc79b76a23..6c3d58a89c 100644
--- a/upbc/message_layout.h
+++ b/upbc/message_layout.h
@@ -116,6 +116,9 @@ class MessageLayout {
   static bool IsPowerOfTwo(size_t val) { return (val & (val - 1)) == 0; }
 
   static size_t Align(size_t val, size_t align) {
+    if (!IsPowerOfTwo(align)) {
+      fprintf(stderr, "YO! Align is: %d\n", (int)align);
+    }
     ABSL_ASSERT(IsPowerOfTwo(align));
     return (val + align - 1) & ~(align - 1);
   }
diff --git a/upbc/protoc-gen-upb.cc b/upbc/protoc-gen-upb.cc
index f282b9c508..5c6ffeba90 100644
--- a/upbc/protoc-gen-upb.cc
+++ b/upbc/protoc-gen-upb.cc
@@ -375,26 +375,34 @@ class FilePlatformLayout {
   FilePlatformLayout(const protobuf::FileDescriptor* fd,
                      upb_MiniTablePlatform platform)
       : platform_(platform) {
-    ComputeLayout(fd);
+    BuildMiniTables(fd);
+    BuildExtensions(fd);
   }
 
-  upb_MiniTable* GetTable(const protobuf::Descriptor* m) const {
+  upb_MiniTable* GetMiniTable(const protobuf::Descriptor* m) const {
     auto it = table_map_.find(m);
     assert(it != table_map_.end());
     return it->second;
   }
 
+  const upb_MiniTable_Extension* GetExtension(
+      const protobuf::FieldDescriptor* fd) const {
+    auto it = extension_map_.find(fd);
+    assert(it != extension_map_.end());
+    return it->second;
+  }
+
  private:
-  void ComputeLayout(const protobuf::FileDescriptor* fd) {
+  void BuildMiniTables(const protobuf::FileDescriptor* fd) {
     for (const auto& m : SortedMessages(fd)) {
       table_map_[m] = MakeMiniTable(m);
     }
     for (const auto& m : SortedMessages(fd)) {
-      upb_MiniTable* mt = GetTable(m);
+      upb_MiniTable* mt = GetMiniTable(m);
       for (const auto& f : FieldNumberOrder(m)) {
         if (f->message_type() && f->message_type()->file() == f->file()) {
           upb_MiniTable_Field* mt_f = FindField(mt, f->number());
-          upb_MiniTable* sub_mt = GetTable(f->message_type());
+          upb_MiniTable* sub_mt = GetMiniTable(f->message_type());
           upb_MiniTable_SetSubMessage(mt, mt_f, sub_mt);
         }
         // TODO: proto2 enums.
@@ -402,6 +410,34 @@ class FilePlatformLayout {
     }
   }
 
+  void BuildExtensions(const protobuf::FileDescriptor* fd) {
+    auto sorted = SortedExtensions(fd);
+    upb::MtDataEncoder e;
+    e.StartMessage(0);
+    for (const auto& f : sorted) {
+      fprintf(stderr, "Field number: %d\n", (int)f->number());
+      e.PutField(static_cast<upb_FieldType>(f->type()), f->number(),
+                 GetFieldModifiers(f));
+    }
+    upb::Status status;
+    size_t ext_count;
+    upb_MiniTable_Extension* exts =
+        upb_MiniTable_BuildExtensions(e.data().data(), e.data().size(),
+                                      &ext_count, arena_.ptr(), status.ptr());
+    if (!exts) {
+      fprintf(stderr, "Error building mini-table: %s\n", status.error_message());
+    }
+    ABSL_ASSERT(exts);
+    if (ext_count != sorted.size()) {
+      fprintf(stderr, "ext_count: %zu, sorted.size(): %zu, data(): %s\n", ext_count,
+              sorted.size(), e.data().c_str());
+    }
+    ABSL_ASSERT(ext_count == sorted.size());
+    for (size_t i = 0; i < ext_count; i++) {
+      extension_map_[sorted[i]] = &exts[i];
+    }
+  }
+
   upb_MiniTable* MakeMiniTable(const protobuf::Descriptor* m) {
     if (m->options().message_set_wire_format()) {
       return upb_MiniTable_BuildMessageSet(platform_, arena_.ptr());
@@ -474,8 +510,11 @@ class FilePlatformLayout {
  private:
   typedef absl::flat_hash_map<const protobuf::Descriptor*, upb_MiniTable*>
       TableMap;
+  typedef absl::flat_hash_map<const protobuf::FieldDescriptor*, upb_MiniTable_Extension*>
+      ExtensionMap;
   upb::Arena arena_;
   TableMap table_map_;
+  ExtensionMap extension_map_;
   upb_MiniTablePlatform platform_;
 };
 
@@ -489,11 +528,16 @@ class FileLayout {
   const protobuf::FileDescriptor* descriptor() const { return descriptor_; }
 
   const upb_MiniTable* GetMiniTable32(const protobuf::Descriptor* m) const {
-    return layout32_.GetTable(m);
+    return layout32_.GetMiniTable(m);
   }
 
   const upb_MiniTable* GetMiniTable64(const protobuf::Descriptor* m) const {
-    return layout64_.GetTable(m);
+    return layout64_.GetMiniTable(m);
+  }
+
+  const upb_MiniTable_Extension* GetExtension(
+      const protobuf::FieldDescriptor* f) const {
+    return layout64_.GetExtension(f);
   }
 
  private: