[NFC] Refactor code to reduce boilerplate for new code annotations

src/ir/metadata.cpp

@@ -130,7 +130,7 @@ bool equal(Function* a, Function* b) {
                  bList.list[i],
                  a->codeAnnotations,
                  b->codeAnnotations,
-                 Function::CodeAnnotation())) {
+                 CodeAnnotation())) {
       return false;
     }
   }

src/parser/contexts.h

@@ -1297,39 +1297,63 @@ struct ParseImplicitTypeDefsCtx : TypeParserCtx<ParseImplicitTypeDefsCtx> {
 };
 
 struct AnnotationParserCtx {
-  // Return the inline hint for a call instruction, if there is one.
-  std::optional<std::uint8_t>
-  getInlineHint(const std::vector<Annotation>& annotations) {
-    // Find and apply (the last) inline hint.
-    const Annotation* hint = nullptr;
+  // Parse annotations into IR.
+  CodeAnnotation parseAnnotations(const std::vector<Annotation>& annotations) {
+    CodeAnnotation ret;
+
+    // Find the hints. For hints with content we must find the last one, which
+    // overrides the others.
+    const Annotation* branchHint = nullptr;
+    const Annotation* inlineHint = nullptr;
     for (auto& a : annotations) {
-      if (a.kind == Annotations::InlineHint) {
-        hint = &a;
+      if (a.kind == Annotations::BranchHint) {
+        branchHint = &a;
+      } else if (a.kind == Annotations::InlineHint) {
+        inlineHint = &a;
       }
     }
-    if (!hint) {
-      return std::nullopt;
-    }
-
-    Lexer lexer(hint->contents);
-    if (lexer.empty()) {
-      std::cerr << "warning: empty InlineHint\n";
-      return std::nullopt;
-    }
 
-    auto str = lexer.takeString();
-    if (!str || str->size() != 1) {
-      std::cerr << "warning: invalid InlineHint string\n";
-      return std::nullopt;
+    // Apply the last branch hint, if valid.
+    if (branchHint) {
+      Lexer lexer(branchHint->contents);
+      if (lexer.empty()) {
+        std::cerr << "warning: empty BranchHint\n";
+      } else {
+        auto str = lexer.takeString();
+        if (!str || str->size() != 1) {
+          std::cerr << "warning: invalid BranchHint string\n";
+        } else {
+          auto value = (*str)[0];
+          if (value != 0 && value != 1) {
+            std::cerr << "warning: invalid BranchHint value\n";
+          } else {
+            ret.branchLikely = bool(value);
+          }
+        }
+      }
     }
 
-    uint8_t value = (*str)[0];
-    if (value > 127) {
-      std::cerr << "warning: invalid InlineHint value\n";
-      return std::nullopt;
+    // Apply the last inline hint, if valid.
+    if (inlineHint) {
+      Lexer lexer(inlineHint->contents);
+      if (lexer.empty()) {
+        std::cerr << "warning: empty InlineHint\n";
+      } else {
+        auto str = lexer.takeString();
+        if (!str || str->size() != 1) {
+          std::cerr << "warning: invalid InlineHint string\n";
+        } else {
+          uint8_t value = (*str)[0];
+          if (value > 127) {
+            std::cerr << "warning: invalid InlineHint value\n";
+          } else {
+            ret.inline_ = value;
+          }
+        }
+      }
     }
 
-    return value;
+    return ret;
   }
 };
 
@@ -2000,10 +2024,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
     if (!type.isSignature()) {
       return in.err(pos, "expected function type");
     }
-    auto likely = getBranchHint(annotations);
-    return withLoc(
-      pos,
-      irBuilder.makeIf(label ? *label : Name{}, type.getSignature(), likely));
+    return withLoc(pos,
+                   irBuilder.makeIf(label ? *label : Name{},
+                                    type.getSignature(),
+                                    parseAnnotations(annotations)));
   }
 
   Result<> visitElse() { return withLoc(irBuilder.visitElse()); }
@@ -2422,8 +2446,8 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                     const std::vector<Annotation>& annotations,
                     Name func,
                     bool isReturn) {
-    auto inline_ = getInlineHint(annotations);
-    return withLoc(pos, irBuilder.makeCall(func, isReturn, inline_));
+    return withLoc(
+      pos, irBuilder.makeCall(func, isReturn, parseAnnotations(annotations)));
   }
 
   Result<> makeCallIndirect(Index pos,
@@ -2433,52 +2457,18 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                             bool isReturn) {
     auto t = getTable(pos, table);
     CHECK_ERR(t);
-    auto inline_ = getInlineHint(annotations);
     return withLoc(pos,
-                   irBuilder.makeCallIndirect(*t, type, isReturn, inline_));
-  }
-
-  // Return the branch hint for a branching instruction, if there is one.
-  std::optional<bool>
-  getBranchHint(const std::vector<Annotation>& annotations) {
-    // Find and apply (the last) branch hint.
-    const Annotation* hint = nullptr;
-    for (auto& a : annotations) {
-      if (a.kind == Annotations::BranchHint) {
-        hint = &a;
-      }
-    }
-    if (!hint) {
-      return std::nullopt;
-    }
-
-    Lexer lexer(hint->contents);
-    if (lexer.empty()) {
-      std::cerr << "warning: empty BranchHint\n";
-      return std::nullopt;
-    }
-
-    auto str = lexer.takeString();
-    if (!str || str->size() != 1) {
-      std::cerr << "warning: invalid BranchHint string\n";
-      return std::nullopt;
-    }
-
-    auto value = (*str)[0];
-    if (value != 0 && value != 1) {
-      std::cerr << "warning: invalid BranchHint value\n";
-      return std::nullopt;
-    }
-
-    return bool(value);
+                   irBuilder.makeCallIndirect(
+                     *t, type, isReturn, parseAnnotations(annotations)));
   }
 
   Result<> makeBreak(Index pos,
                      const std::vector<Annotation>& annotations,
                      Index label,
                      bool isConditional) {
-    auto likely = getBranchHint(annotations);
-    return withLoc(pos, irBuilder.makeBreak(label, isConditional, likely));
+    return withLoc(
+      pos,
+      irBuilder.makeBreak(label, isConditional, parseAnnotations(annotations)));
   }
 
   Result<> makeSwitch(Index pos,
@@ -2617,8 +2607,9 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                        const std::vector<Annotation>& annotations,
                        HeapType type,
                        bool isReturn) {
-    auto inline_ = getInlineHint(annotations);
-    return withLoc(pos, irBuilder.makeCallRef(type, isReturn, inline_));
+    return withLoc(
+      pos,
+      irBuilder.makeCallRef(type, isReturn, parseAnnotations(annotations)));
   }
 
   Result<> makeRefI31(Index pos,
@@ -2658,8 +2649,9 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                     BrOnOp op,
                     Type in = Type::none,
                     Type out = Type::none) {
-    auto likely = getBranchHint(annotations);
-    return withLoc(pos, irBuilder.makeBrOn(label, op, in, out, likely));
+    return withLoc(
+      pos,
+      irBuilder.makeBrOn(label, op, in, out, parseAnnotations(annotations)));
   }
 
   Result<> makeStructNew(Index pos,

src/wasm-binary.h

@@ -1719,17 +1719,19 @@ class WasmBinaryReader {
   void readDylink(size_t payloadLen);
   void readDylink0(size_t payloadLen);
 
-  // We read branch hints *after* the code section, even though they appear
+  // We read code annotations *after* the code section, even though they appear
   // earlier. That is simpler for us as we note expression locations as we scan
-  // code, and then just need to match them up. To do this, we note the branch
-  // 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);
+  // code, and then just need to match them up. To do this, we note the
+  // positions of annotation sections in the first pass, and handle them later.
+  struct AnnotationSectionInfo {
+    // The start position of the section. We will rewind to there to read it.
+    size_t pos;
+    // A lambda that will read the section, from that position.
+    std::function<void()> read;
+  };
+  std::vector<AnnotationSectionInfo> deferredAnnotationSections;
 
-  // Like branch hints, we note where the section is to read it later.
-  size_t inlineHintsPos = 0;
-  size_t inlineHintsLen = 0;
+  void readBranchHints(size_t payloadLen);
   void readInlineHints(size_t payloadLen);
 
   std::tuple<Address, Address, Index, MemoryOrder>

src/wasm-ir-builder.h

@@ -129,20 +129,19 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeNop();
   Result<> makeBlock(Name label, Signature sig);
   Result<>
-  makeIf(Name label, Signature sig, std::optional<bool> likely = std::nullopt);
+  makeIf(Name label, Signature sig, const CodeAnnotation& annotations = {});
   Result<> makeLoop(Name label, Signature sig);
   Result<> makeBreak(Index label,
                      bool isConditional,
-                     std::optional<bool> likely = std::nullopt);
+                     const CodeAnnotation& annotations = {});
   Result<> makeSwitch(const std::vector<Index>& labels, Index defaultLabel);
   // Unlike Builder::makeCall, this assumes the function already exists.
-  Result<> makeCall(Name func,
-                    bool isReturn,
-                    std::optional<std::uint8_t> inline_ = std::nullopt);
+  Result<>
+  makeCall(Name func, bool isReturn, const CodeAnnotation& annotations = {});
   Result<> makeCallIndirect(Name table,
                             HeapType type,
                             bool isReturn,
-                            std::optional<std::uint8_t> inline_ = std::nullopt);
+                            const CodeAnnotation& annotations = {});
   Result<> makeLocalGet(Index local);
   Result<> makeLocalSet(Index local);
   Result<> makeLocalTee(Index local);
@@ -226,15 +225,15 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeI31Get(bool signed_);
   Result<> makeCallRef(HeapType type,
                        bool isReturn,
-                       std::optional<std::uint8_t> inline_ = std::nullopt);
+                       const CodeAnnotation& annotations = {});
   Result<> makeRefTest(Type type);
   Result<> makeRefCast(Type type, bool isDesc);
   Result<> makeRefGetDesc(HeapType type);
   Result<> makeBrOn(Index label,
                     BrOnOp op,
                     Type in = Type::none,
                     Type out = Type::none,
-                    std::optional<bool> likely = std::nullopt);
+                    const CodeAnnotation& annotations = {});
   Result<> makeStructNew(HeapType type, bool isDesc);
   Result<> makeStructNewDefault(HeapType type, bool isDesc);
   Result<>
@@ -735,11 +734,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Expression* fixExtraOutput(ScopeCtx& scope, Name label, Expression* expr);
   void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);
 
-  // Add a branch hint, if |likely| is present.
-  void addBranchHint(Expression* expr, std::optional<bool> likely);
-
-  // Add an inlining hint, if |inline_| is present.
-  void addInlineHint(Expression* expr, std::optional<std::uint8_t> inline_);
+  void applyAnnotations(Expression* expr, const CodeAnnotation& annotation);
 
   void dump();
 };

src/wasm.h

@@ -2232,6 +2232,21 @@ struct BinaryLocations {
 // Forward declaration for FuncEffectsMap.
 class EffectAnalyzer;
 
+// Code annotations for VMs.
+struct CodeAnnotation {
+  // Branch Hinting proposal: Whether the branch is likely, or unlikely.
+  std::optional<bool> branchLikely;
+
+  // Compilation Hints proposal.
+  static const uint8_t NeverInline = 0;
+  static const uint8_t AlwaysInline = 127;
+  std::optional<uint8_t> inline_;
+
+  bool operator==(const CodeAnnotation& other) const {
+    return branchLikely == other.branchLikely && inline_ == other.inline_;
+  }
+};
+
 class Function : public Importable {
 public:
   // A non-nullable reference to a function type. Exact for defined functions.
@@ -2285,25 +2300,10 @@ class Function : public Importable {
     delimiterLocations;
   BinaryLocations::FunctionLocations funcLocation;
 
-  // Code annotations for VMs. As with debug info, we do not store these on
+  // Function-level annotations are implemented with a key of nullptr, matching
+  // the 0 byte offset in the spec. As with debug info, we do not store these on
   // Expressions as we assume most instances are unannotated, and do not want to
   // add constant memory overhead.
-  struct CodeAnnotation {
-    // Branch Hinting proposal: Whether the branch is likely, or unlikely.
-    std::optional<bool> branchLikely;
-
-    // Compilation Hints proposal.
-    static const uint8_t NeverInline = 0;
-    static const uint8_t AlwaysInline = 127;
-    std::optional<uint8_t> inline_;
-
-    bool operator==(const CodeAnnotation& other) const {
-      return branchLikely == other.branchLikely && inline_ == other.inline_;
-    }
-  };
-
-  // Function-level annotations are implemented with a key of nullptr, matching
-  // the 0 byte offset in the spec.
   std::unordered_map<Expression*, CodeAnnotation> codeAnnotations;
 
   // The effects for this function, if they have been computed. We use a shared