[Code Annotations] Start Branch Hinting support, in the text format, for Breaks (#7567)

This adds the first minimal amount of work for Branch Hinting (and
code annotations, more generally): Just the text format, and just a
single instruction (Break).

This design follows debug info: Assuming most instructions have no
branch hints (like most do not have debug info), we do not add a field
to the Break class itself, but to metadata on the side.

---------

Co-authored-by: Thomas Lively <[email protected]>

Author: kripken

src/parser/contexts.h

@@ -23,6 +23,7 @@
 #include "support/name.h"
 #include "support/result.h"
 #include "support/string.h"
+#include "wasm-annotations.h"
 #include "wasm-builder.h"
 #include "wasm-ir-builder.h"
 #include "wasm.h"
@@ -2339,11 +2340,47 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return withLoc(pos, irBuilder.makeCallIndirect(*t, type, isReturn));
   }
 
+  // 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);
+  }
+
   Result<> makeBreak(Index pos,
                      const std::vector<Annotation>& annotations,
                      Index label,
                      bool isConditional) {
-    return withLoc(pos, irBuilder.makeBreak(label, isConditional));
+    auto likely = getBranchHint(annotations);
+    return withLoc(pos, irBuilder.makeBreak(label, isConditional, likely));
   }
 
   Result<> makeSwitch(Index pos,

src/passes/Print.cpp

@@ -27,6 +27,7 @@
 #include <pass.h>
 #include <pretty_printing.h>
 #include <support/string.h>
+#include <wasm-annotations.h>
 #include <wasm-stack.h>
 #include <wasm-type-printing.h>
 #include <wasm.h>
@@ -267,15 +268,17 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
 
   void
   printDebugLocation(const std::optional<Function::DebugLocation>& location);
-  void printDebugLocation(Expression* curr);
 
   // Prints debug info for a delimiter in an expression.
   void printDebugDelimiterLocation(Expression* curr, Index i);
 
+  // Prints debug info and code annotations.
+  void printMetadata(Expression* curr);
+
   void printExpressionContents(Expression* curr);
 
   void visit(Expression* curr) {
-    printDebugLocation(curr);
+    printMetadata(curr);
     UnifiedExpressionVisitor<PrintSExpression>::visit(curr);
   }
 
@@ -2677,27 +2680,40 @@ void PrintSExpression::printDebugLocation(
   doIndent(o, indent);
 }
 
-void PrintSExpression::printDebugLocation(Expression* curr) {
+void PrintSExpression::printMetadata(Expression* curr) {
   if (currFunction) {
-    // show an annotation, if there is one
-    auto& debugLocations = currFunction->debugLocations;
-    auto iter = debugLocations.find(curr);
-    if (iter != debugLocations.end()) {
+    // Show a debug location, if there is one.
+    if (auto iter = currFunction->debugLocations.find(curr);
+        iter != currFunction->debugLocations.end()) {
       printDebugLocation(iter->second);
     } else {
       printDebugLocation(std::nullopt);
     }
-    // show a binary position, if there is one
+
+    // Show a binary position, if there is one.
     if (debugInfo) {
-      auto iter = currFunction->expressionLocations.find(curr);
-      if (iter != currFunction->expressionLocations.end()) {
+      if (auto iter = currFunction->expressionLocations.find(curr);
+          iter != currFunction->expressionLocations.end()) {
         Colors::grey(o);
         o << ";; code offset: 0x" << std::hex << iter->second.start << std::dec
           << '\n';
         restoreNormalColor(o);
         doIndent(o, indent);
       }
     }
+
+    // Show a code annotation, if there is one.
+    if (auto iter = currFunction->codeAnnotations.find(curr);
+        iter != currFunction->codeAnnotations.end()) {
+      auto& annotation = iter->second;
+      if (annotation.branchLikely) {
+        Colors::grey(o);
+        o << "(@" << Annotations::BranchHint << " \"\\0"
+          << (*annotation.branchLikely ? "1" : "0") << "\")\n";
+        restoreNormalColor(o);
+        doIndent(o, indent);
+      }
+    }
   }
 }
 
@@ -2781,7 +2797,7 @@ void PrintSExpression::visitBlock(Block* curr) {
   while (1) {
     if (stack.size() > 0) {
       doIndent(o, indent);
-      printDebugLocation(curr);
+      printMetadata(curr);
     }
     stack.push_back(curr);
     o << '(';

src/wasm-annotations.h

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Support for code annotations.
+//
+
+#ifndef wasm_annotations_h
+#define wasm_annotations_h
+
+#include "support/name.h"
+
+namespace wasm::Annotations {
+
+extern const Name BranchHint;
+
+} // namespace wasm::Annotations
+
+#endif // wasm_annotations_h

src/wasm-ir-builder.h

@@ -72,8 +72,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     this->codeSectionOffset = codeSectionOffset;
   }
 
-  // Set the function used to add scratch locals when constructing an isolated
-  // sequence of IR.
+  // The current function, used to create scratch locals, add annotations, etc.
   void setFunction(Function* func) { this->func = func; }
 
   // Handle the boundaries of control flow structures. Users may choose to use
@@ -119,7 +118,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeBlock(Name label, Signature sig);
   Result<> makeIf(Name label, Signature sig);
   Result<> makeLoop(Name label, Signature sig);
-  Result<> makeBreak(Index label, bool isConditional);
+  Result<> makeBreak(Index label,
+                     bool isConditional,
+                     std::optional<bool> likely = std::nullopt);
   Result<> makeSwitch(const std::vector<Index>& labels, Index defaultLabel);
   // Unlike Builder::makeCall, this assumes the function already exists.
   Result<> makeCall(Name func, bool isReturn);

src/wasm.h

@@ -2191,6 +2191,16 @@ class Function : public Importable {
     delimiterLocations;
   BinaryLocations::FunctionLocations funcLocation;
 
+  // Code annotations. As with debug info, we do not store these on Expressions
+  // themselves, 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;
+  };
+
+  std::unordered_map<Expression*, CodeAnnotation> codeAnnotations;
+
   // The effects for this function, if they have been computed. We use a shared
   // ptr here to avoid compilation errors with the forward-declared
   // EffectAnalyzer.

src/wasm/wasm-ir-builder.cpp

@@ -1393,7 +1393,9 @@ Result<> IRBuilder::makeLoop(Name label, Signature sig) {
   return visitLoopStart(loop, sig.params);
 }
 
-Result<> IRBuilder::makeBreak(Index label, bool isConditional) {
+Result<> IRBuilder::makeBreak(Index label,
+                              bool isConditional,
+                              std::optional<bool> likely) {
   auto name = getLabelName(label);
   CHECK_ERR(name);
   auto labelType = getLabelType(label);
@@ -1404,7 +1406,15 @@ Result<> IRBuilder::makeBreak(Index label, bool isConditional) {
   // Use a dummy condition value if we need to pop a condition.
   curr.condition = isConditional ? &curr : nullptr;
   CHECK_ERR(ChildPopper{*this}.visitBreak(&curr, *labelType));
-  push(builder.makeBreak(curr.name, curr.value, curr.condition));
+  auto* br = builder.makeBreak(curr.name, curr.value, curr.condition);
+  push(br);
+
+  if (likely) {
+    // Branches are only possible inside functions.
+    assert(func);
+    func->codeAnnotations[br].branchLikely = likely;
+  }
+
   return Ok{};
 }
 

src/wasm/wasm.cpp

@@ -16,6 +16,7 @@
 
 #include "wasm.h"
 #include "ir/branch-utils.h"
+#include "wasm-annotations.h"
 #include "wasm-traversal.h"
 #include "wasm-type.h"
 
@@ -63,6 +64,12 @@ const char* CustomDescriptorsFeature = "custom-descriptors";
 
 } // namespace BinaryConsts::CustomSections
 
+namespace Annotations {
+
+const Name BranchHint = "metadata.code.branch_hint";
+
+} // namespace Annotations
+
 Name STACK_POINTER("__stack_pointer");
 Name MODULE("module");
 Name START("start");

test/lit/wat-annotations.wast

@@ -0,0 +1,43 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt -all %s -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $0 (func (param i32)))
+
+  ;; CHECK:      (func $branch-hints-br_if (type $0) (param $x i32)
+  ;; CHECK-NEXT:  (block $out
+  ;; CHECK-NEXT:   (@metadata.code.branch_hint "\00")
+  ;; CHECK-NEXT:   (br_if $out
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (@metadata.code.branch_hint "\01")
+  ;; CHECK-NEXT:   (br_if $out
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (@metadata.code.branch_hint "\00")
+  ;; CHECK-NEXT:   (br_if $out
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $branch-hints-br_if (param $x i32)
+    (block $out
+      ;; A branch annotated as unlikely, and one as likely.
+      (@metadata.code.branch_hint "\00")
+      (br_if $out
+        (local.get $x)
+      )
+      (@metadata.code.branch_hint "\01")
+      (br_if $out
+        (local.get $x)
+      )
+      ;; The last one wins.
+      (@metadata.code.branch_hint "\01")
+      (@metadata.code.branch_hint "\00")
+      (br_if $out
+        (local.get $x)
+      )
+    )
+  )
+)