Implement custom section reading/writing (WebAssembly#2284)

include/wabt/binary-reader-logging.h

@@ -346,6 +346,12 @@ class BinaryReaderLogging : public BinaryReaderDelegate {
   Result OnDylinkExport(std::string_view name, uint32_t flags) override;
   Result EndDylinkSection() override;
 
+  Result BeginGenericCustomSection(Offset size) override;
+  Result OnGenericCustomSection(std::string_view name,
+                                const void* data,
+                                Offset size) override;
+  Result EndGenericCustomSection() override;
+
   Result BeginTargetFeaturesSection(Offset size) override;
   Result OnFeatureCount(Index count) override;
   Result OnFeature(uint8_t prefix, std::string_view name) override;

include/wabt/binary-reader-nop.h

@@ -519,6 +519,15 @@ class BinaryReaderNop : public BinaryReaderDelegate {
   }
   Result EndTargetFeaturesSection() override { return Result::Ok; }
 
+  /* Generic custom section */
+  Result BeginGenericCustomSection(Offset size) override { return Result::Ok; }
+  Result OnGenericCustomSection(std::string_view name,
+                                const void* data,
+                                Offset size) override {
+    return Result::Ok;
+  };
+  Result EndGenericCustomSection() override { return Result::Ok; }
+
   /* Linking section */
   Result BeginLinkingSection(Offset size) override { return Result::Ok; }
   Result OnSymbolCount(Index count) override { return Result::Ok; }

include/wabt/binary-reader.h

@@ -424,6 +424,13 @@ class BinaryReaderDelegate {
   virtual Result OnFeature(uint8_t prefix, std::string_view name) = 0;
   virtual Result EndTargetFeaturesSection() = 0;
 
+  /* Generic custom section */
+  virtual Result BeginGenericCustomSection(Offset size) = 0;
+  virtual Result OnGenericCustomSection(std::string_view name,
+                                        const void* data,
+                                        Offset size) = 0;
+  virtual Result EndGenericCustomSection() = 0;
+
   /* Linking section */
   virtual Result BeginLinkingSection(Offset size) = 0;
   virtual Result OnSymbolCount(Index count) = 0;

include/wabt/ir.h

@@ -1164,6 +1164,17 @@ class StartModuleField : public ModuleFieldMixin<ModuleFieldType::Start> {
   Var start;
 };
 
+struct Custom {
+  explicit Custom(const Location& loc = Location(),
+                  std::string_view name = std::string_view(),
+                  std::vector<uint8_t> data = std::vector<uint8_t>())
+      : name(name), data(data), loc(loc) {}
+
+  std::string name;
+  std::vector<uint8_t> data;
+  Location loc;
+};
+
 struct Module {
   Index GetFuncTypeIndex(const Var&) const;
   Index GetFuncTypeIndex(const FuncDeclaration&) const;
@@ -1235,6 +1246,7 @@ struct Module {
   std::vector<Memory*> memories;
   std::vector<DataSegment*> data_segments;
   std::vector<Var*> starts;
+  std::vector<Custom> customs;
 
   BindingHash tag_bindings;
   BindingHash func_bindings;

include/wabt/token.def

@@ -20,6 +20,7 @@
 
 /* Tokens with no additional data (i.e. bare). */
 WABT_TOKEN(Invalid, "Invalid")
+WABT_TOKEN(After, "after")
 WABT_TOKEN(Array, "array")
 WABT_TOKEN(AssertException, "assert_exception")
 WABT_TOKEN(AssertExhaustion, "assert_exhaustion")
@@ -28,6 +29,7 @@ WABT_TOKEN(AssertMalformed, "assert_malformed")
 WABT_TOKEN(AssertReturn, "assert_return")
 WABT_TOKEN(AssertTrap, "assert_trap")
 WABT_TOKEN(AssertUnlinkable, "assert_unlinkable")
+WABT_TOKEN(Before, "before")
 WABT_TOKEN(Bin, "bin")
 WABT_TOKEN(Item, "item")
 WABT_TOKEN(Data, "data")
@@ -40,6 +42,7 @@ WABT_TOKEN(Eof, "EOF")
 WABT_TOKEN(Tag, "tag")
 WABT_TOKEN(Export, "export")
 WABT_TOKEN(Field, "field")
+WABT_TOKEN(Function, "function")
 WABT_TOKEN(Get, "get")
 WABT_TOKEN(Global, "global")
 WABT_TOKEN(Import, "import")
@@ -95,6 +98,7 @@ WABT_TOKEN(Block, "block")
 WABT_TOKEN(Br, "br")
 WABT_TOKEN(BrIf, "br_if")
 WABT_TOKEN(BrTable, "br_table")
+WABT_TOKEN(Code, "code")
 WABT_TOKEN(Call, "call")
 WABT_TOKEN(CallIndirect, "call_indirect")
 WABT_TOKEN(CallRef, "call_ref")

include/wabt/wast-parser.h

@@ -163,6 +163,9 @@ class WastParser {
   Result ParseStartModuleField(Module*);
   Result ParseTableModuleField(Module*);
 
+  Result ParseCustomSectionAnnotation(Module*);
+  bool PeekIsCustom();
+
   Result ParseExportDesc(Export*);
   Result ParseInlineExports(ModuleFieldList*, ExternalKind);
   Result ParseInlineImport(Import*);

src/binary-reader-ir.cc

@@ -314,6 +314,10 @@ class BinaryReaderIR : public BinaryReaderNop {
                      Index index,
                      std::string_view name) override;
 
+  Result OnGenericCustomSection(std::string_view name,
+                                const void* data,
+                                Offset size) override;
+
   Result BeginTagSection(Offset size) override { return Result::Ok; }
   Result OnTagCount(Index count) override { return Result::Ok; }
   Result OnTagType(Index index, Index sig_index) override;
@@ -1750,6 +1754,18 @@ Result BinaryReaderIR::OnTableSymbol(Index index,
   return SetTableName(table_index, name);
 }
 
+Result BinaryReaderIR::OnGenericCustomSection(std::string_view name,
+                                              const void* data,
+                                              Offset size) {
+  Custom custom = Custom(GetLocation(), name);
+  custom.data.resize(size);
+  if (size > 0) {
+    memcpy(custom.data.data(), data, size);
+  }
+  module_->customs.push_back(std::move(custom));
+  return Result::Ok;
+}
+
 }  // end anonymous namespace
 
 Result ReadBinaryIr(const char* filename,

src/binary-reader-logging.cc

@@ -654,6 +654,15 @@ Result BinaryReaderLogging::OnCodeMetadata(Offset code_offset,
   return reader_->OnCodeMetadata(code_offset, data, size);
 }
 
+Result BinaryReaderLogging::OnGenericCustomSection(std::string_view name,
+                                                   const void* data,
+                                                   Offset size) {
+  LOGF("OnGenericCustomSection(name: \"" PRIstringview "\", size: %" PRIzd
+       ")\n",
+       WABT_PRINTF_STRING_VIEW_ARG(name), size);
+  return reader_->OnGenericCustomSection(name, data, size);
+}
+
 #define DEFINE_BEGIN(name)                        \
   Result BinaryReaderLogging::name(Offset size) { \
     LOGF(#name "(%" PRIzd ")\n", size);           \
@@ -894,6 +903,9 @@ DEFINE_INDEX(OnInitFunctionCount)
 DEFINE_INDEX(OnComdatCount)
 DEFINE_END(EndLinkingSection)
 
+DEFINE_BEGIN(BeginGenericCustomSection);
+DEFINE_END(EndGenericCustomSection);
+
 DEFINE_BEGIN(BeginTagSection);
 DEFINE_INDEX(OnTagCount);
 DEFINE_INDEX_INDEX(OnTagType, "index", "sig_index")

src/binary-reader.cc

@@ -116,6 +116,9 @@ class BinaryReader {
   [[nodiscard]] Result ReadBytes(const void** out_data,
                                  Address* out_data_size,
                                  const char* desc);
+  [[nodiscard]] Result ReadBytesWithSize(const void** out_data,
+                                         Offset size,
+                                         const char* desc);
   [[nodiscard]] Result ReadIndex(Index* index, const char* desc);
   [[nodiscard]] Result ReadOffset(Offset* offset, const char* desc);
   [[nodiscard]] Result ReadAlignment(Address* align_log2, const char* desc);
@@ -153,6 +156,8 @@ class BinaryReader {
   [[nodiscard]] Result ReadNameSection(Offset section_size);
   [[nodiscard]] Result ReadRelocSection(Offset section_size);
   [[nodiscard]] Result ReadDylinkSection(Offset section_size);
+  [[nodiscard]] Result ReadGenericCustomSection(std::string_view name,
+                                                Offset section_size);
   [[nodiscard]] Result ReadDylink0Section(Offset section_size);
   [[nodiscard]] Result ReadTargetFeaturesSections(Offset section_size);
   [[nodiscard]] Result ReadLinkingSection(Offset section_size);
@@ -394,13 +399,19 @@ Result BinaryReader::ReadBytes(const void** out_data,
                                const char* desc) {
   uint32_t data_size = 0;
   CHECK_RESULT(ReadU32Leb128(&data_size, "data size"));
+  CHECK_RESULT(ReadBytesWithSize(out_data, data_size, desc));
+  *out_data_size = data_size;
+  return Result::Ok;
+}
 
-  ERROR_UNLESS(state_.offset + data_size <= read_end_,
-               "unable to read data: %s", desc);
+Result BinaryReader::ReadBytesWithSize(const void** out_data,
+                                       Offset size,
+                                       const char* desc) {
+  ERROR_UNLESS(state_.offset + size <= read_end_, "unable to read data: %s",
+               desc);
 
   *out_data = static_cast<const uint8_t*>(state_.data) + state_.offset;
-  *out_data_size = data_size;
-  state_.offset += data_size;
+  state_.offset += size;
   return Result::Ok;
 }
 
@@ -2155,6 +2166,18 @@ Result BinaryReader::ReadTargetFeaturesSections(Offset section_size) {
   return Result::Ok;
 }
 
+Result BinaryReader::ReadGenericCustomSection(std::string_view name,
+                                              Offset section_size) {
+  CALLBACK(BeginGenericCustomSection, section_size);
+  const void* data;
+  Offset custom_data_size = read_end_ - state_.offset;
+  CHECK_RESULT(
+      ReadBytesWithSize(&data, custom_data_size, "custom section data"));
+  CALLBACK(OnGenericCustomSection, name, data, custom_data_size);
+  CALLBACK0(EndGenericCustomSection);
+  return Result::Ok;
+}
+
 Result BinaryReader::ReadLinkingSection(Offset section_size) {
   CALLBACK(BeginLinkingSection, section_size);
   uint32_t version;
@@ -2375,6 +2398,13 @@ Result BinaryReader::ReadCustomSection(Index section_index,
   ValueRestoreGuard<bool, &BinaryReader::reading_custom_section_> guard(this);
   reading_custom_section_ = true;
 
+  {
+    // Backtrack parser when scope ends
+    ValueRestoreGuard<BinaryReaderDelegate::State, &BinaryReader::state_> guard(
+        this);
+    CHECK_RESULT(ReadGenericCustomSection(section_name, section_size));
+  }
+
   if (options_.read_debug_names && section_name == WABT_BINARY_SECTION_NAME) {
     CHECK_RESULT(ReadNameSection(section_size));
     did_read_names_section_ = true;
@@ -2395,7 +2425,7 @@ Result BinaryReader::ReadCustomSection(Index section_index,
     metadata_name.remove_prefix(sizeof(WABT_BINARY_SECTION_CODE_METADATA) - 1);
     CHECK_RESULT(ReadCodeMetadataSection(metadata_name, section_size));
   } else {
-    // This is an unknown custom section, skip it.
+    // Skip. This is a generic custom section, and is handled above.
     state_.offset = read_end_;
   }
   CALLBACK0(EndCustomSection);

src/binary-writer.cc

@@ -1670,6 +1670,23 @@ Result BinaryWriter::WriteModule() {
     EndSection();
   }
 
+  for (const Custom& custom : module_->customs) {
+    // These custom sections are already specially handled by BinaryWriter, so
+    // we don't want to double-write.
+    if ((custom.name == WABT_BINARY_SECTION_NAME &&
+         options_.write_debug_names) ||
+        (custom.name.rfind(WABT_BINARY_SECTION_RELOC) == 0 &&
+         options_.relocatable) ||
+        (custom.name == WABT_BINARY_SECTION_LINKING && options_.relocatable) ||
+        (custom.name.find(WABT_BINARY_SECTION_CODE_METADATA) == 0 &&
+         options_.features.code_metadata_enabled())) {
+      continue;
+    }
+    BeginCustomSection(custom.name.data());
+    stream_->WriteData(custom.data, "custom data");
+    EndSection();
+  }
+
   if (options_.write_debug_names) {
     std::vector<std::string> index_to_name;
 

src/lexer-keywords.txt

@@ -18,6 +18,7 @@ struct TokenInfo {
 };
 %%
 array, Type::Array, TokenType::Array
+after, TokenType::After
 assert_exception, TokenType::AssertException
 assert_exhaustion, TokenType::AssertExhaustion
 assert_invalid, TokenType::AssertInvalid
@@ -26,6 +27,7 @@ assert_return, TokenType::AssertReturn
 assert_trap, TokenType::AssertTrap
 assert_unlinkable, TokenType::AssertUnlinkable
 atomic.fence, TokenType::AtomicFence, Opcode::AtomicFence
+before, TokenType::Before
 binary, TokenType::Bin
 block, TokenType::Block, Opcode::Block
 br_if, TokenType::BrIf, Opcode::BrIf
@@ -36,6 +38,7 @@ call_ref, TokenType::CallRef, Opcode::CallRef
 call, TokenType::Call, Opcode::Call
 catch, TokenType::Catch, Opcode::Catch
 catch_all, TokenType::CatchAll, Opcode::CatchAll
+code, TokenType::Code
 data.drop, TokenType::DataDrop, Opcode::DataDrop
 data, TokenType::Data
 declare, TokenType::Declare
@@ -178,6 +181,7 @@ f64x2, TokenType::F64X2
 field, TokenType::Field
 funcref, Type::FuncRef
 func, Type::FuncRef, TokenType::Func
+function, TokenType::Function
 get, TokenType::Get
 global.get, TokenType::GlobalGet, Opcode::GlobalGet
 global.set, TokenType::GlobalSet, Opcode::GlobalSet