[NFC] Refactor a generic function for binary writing of a code annotations section (#7592)

This will allow 95% of the code to be shared for binary writing, between
branch hints and inline hints for example. They have almost identical
formats in structure, differing only in the hint itself and section
name.

Also prepare for multiple code annotations sections.

Author: kripken

src/wasm-binary.h

@@ -1410,6 +1410,8 @@ class WasmBinaryWriter {
   // must then insert before the code (as the spec requires that).
   std::optional<BufferWithRandomAccess> writeCodeAnnotations();
 
+  std::optional<BufferWithRandomAccess> getBranchHintsBuffer();
+
   // helpers
   void writeInlineString(std::string_view name);
   void writeEscapedName(std::string_view name);
@@ -1492,6 +1494,15 @@ class WasmBinaryWriter {
   std::unordered_map<Name, Index> stringIndexes;
 
   void prepare();
+
+  // Internal helper for emitting a code annotation section for a hint that is
+  // expression offset based. Receives the name of the section and two
+  // functions, one to check if the annotation we care about exists (receiving
+  // the annotation), and another to emit it (receiving the annotation and the
+  // buffer to write in).
+  template<typename HasFunc, typename EmitFunc>
+  std::optional<BufferWithRandomAccess>
+  writeExpressionHints(Name sectionName, HasFunc has, EmitFunc emit);
 };
 
 extern std::vector<char> defaultEmptySourceMap;
@@ -1679,7 +1690,6 @@ class WasmBinaryReader {
   // hint position and size in the first pass, and handle it later.
   size_t branchHintsPos = 0;
   size_t branchHintsLen = 0;
-
   void readBranchHints(size_t payloadLen);
 
   Index readMemoryAccess(Address& alignment, Address& offset);

src/wasm/wasm-binary.cpp

@@ -1549,8 +1549,28 @@ void WasmBinaryWriter::trackExpressionDelimiter(Expression* curr,
 }
 
 std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeCodeAnnotations() {
-  // Assemble the info for Branch Hinting: for each function, a vector of the
-  // hints.
+  std::optional<BufferWithRandomAccess> ret;
+
+  auto append = [&](std::optional<BufferWithRandomAccess>&& temp) {
+    if (temp) {
+      if (!ret) {
+        // This is the first section.
+        ret = std::move(temp);
+      } else {
+        // This is a later section, append.
+        ret->insert(ret->end(), temp->begin(), temp->end());
+      }
+    }
+  };
+
+  append(getBranchHintsBuffer());
+  return ret;
+}
+
+template<typename HasFunc, typename EmitFunc>
+std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeExpressionHints(
+  Name sectionName, HasFunc has, EmitFunc emit) {
+  // Assemble the info: for each function, a vector of the hints.
   struct ExprHint {
     Expression* expr;
     // The offset we will write in the custom section.
@@ -1566,15 +1586,15 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeCodeAnnotations() {
   std::vector<FuncHints> funcHintsVec;
 
   for (auto& func : wasm->functions) {
-    // Collect the Branch Hints for this function.
+    // Collect the hints for this function.
     FuncHints funcHints;
 
     // We compute the location of the function declaration area (where the
     // locals are declared) the first time we need it.
     BinaryLocation funcDeclarationsOffset = 0;
 
     for (auto& [expr, annotation] : func->codeAnnotations) {
-      if (annotation.branchLikely) {
+      if (has(annotation)) {
         auto exprIter = binaryLocations.expressions.find(expr);
         if (exprIter == binaryLocations.expressions.end()) {
           // No expression exists for this annotation - perhaps optimizations
@@ -1622,7 +1642,7 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeCodeAnnotations() {
   // We found data: emit the section.
   buffer << uint8_t(BinaryConsts::Custom);
   auto lebPos = buffer.writeU32LEBPlaceholder();
-  buffer.writeInlineString(Annotations::BranchHint.str);
+  buffer.writeInlineString(sectionName.str);
 
   buffer << U32LEB(funcHintsVec.size());
   for (auto& funcHints : funcHintsVec) {
@@ -1632,14 +1652,7 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeCodeAnnotations() {
     for (auto& exprHint : funcHints.exprHints) {
       buffer << U32LEB(exprHint.offset);
 
-      // Hint size, always 1 for now.
-      buffer << U32LEB(1);
-
-      // We must only emit hints that are present.
-      assert(exprHint.hint->branchLikely);
-
-      // Hint contents: likely or not.
-      buffer << U32LEB(int(*exprHint.hint->branchLikely));
+      emit(*exprHint.hint, buffer);
     }
   }
 
@@ -1651,6 +1664,25 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeCodeAnnotations() {
   return buffer;
 }
 
+std::optional<BufferWithRandomAccess> WasmBinaryWriter::getBranchHintsBuffer() {
+  return writeExpressionHints(
+    Annotations::BranchHint,
+    [](const Function::CodeAnnotation& annotation) {
+      return annotation.branchLikely;
+    },
+    [](const Function::CodeAnnotation& annotation,
+       BufferWithRandomAccess& buffer) {
+      // Hint size, always 1 for now.
+      buffer << U32LEB(1);
+
+      // We must only emit hints that are present.
+      assert(annotation.branchLikely);
+
+      // Hint contents: likely or not.
+      buffer << U32LEB(int(*annotation.branchLikely));
+    });
+}
+
 void WasmBinaryWriter::writeData(const char* data, size_t size) {
   for (size_t i = 0; i < size; i++) {
     o << int8_t(data[i]);