[CodeAnnotations] Implement function-level inlining hints

src/parser/contexts.h

@@ -1474,13 +1474,11 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
         Builder::addVar(f.get(), l.name, l.type);
       }
     }
-    // TODO: Add function-level annotations (stored using the nullptr key, as
-    // they are tied to no instruction in particular), but this should wait on
-    // figuring out
-    // https://github.com/WebAssembly/tool-conventions/issues/251
-    // if (auto inline_ = getInlineHint(annotations)) {
-    //   f->codeAnnotations[nullptr].inline_ = inline_;
-    // }
+    // Function-level annotations are stored using the nullptr key, as they are
+    // not tied to a particular instruction.
+    if (!annotations.empty()) {
+      f->codeAnnotations[nullptr] = parseAnnotations(annotations);
+    }
     return Ok{};
   }
 

src/passes/Print.cpp

@@ -3246,6 +3246,7 @@ void PrintSExpression::visitImportedFunction(Function* curr) {
   lastPrintedLocation = std::nullopt;
   o << '(';
   emitImportHeader(curr);
+  printCodeAnnotations(nullptr);
   handleSignature(curr);
   o << "))";
   o << maybeNewLine;
@@ -3259,9 +3260,7 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
   if (currFunction->prologLocation) {
     printDebugLocation(*currFunction->prologLocation);
   }
-  // TODO: print code annotations in the right place, depending on
-  // https://github.com/WebAssembly/tool-conventions/issues/251
-  // printCodeAnnotations(nullptr);
+  printCodeAnnotations(nullptr);
   handleSignature(curr, true);
   incIndent();
   for (size_t i = curr->getVarIndexBase(); i < curr->getNumLocals(); i++) {

src/wasm/wasm-binary.cpp

@@ -456,7 +456,11 @@ void WasmBinaryWriter::writeFunctions() {
         }
       }
     }
-    if (!binaryLocationTrackedExpressionsForFunc.empty()) {
+    // We need to track the function location if we are tracking the locations
+    // of expressions inside it, or, if it has code annotations (the function
+    // itself may be annotated, even if nothing inside it is).
+    if (!binaryLocationTrackedExpressionsForFunc.empty() ||
+        !func->codeAnnotations.empty()) {
       binaryLocations.functions[func] = BinaryLocations::FunctionLocations{
         BinaryLocation(sizePos),
         BinaryLocation(start - adjustmentForLEBShrinking),
@@ -1654,23 +1658,30 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeExpressionHints(
 
     for (auto& [expr, annotation] : func->codeAnnotations) {
       if (has(annotation)) {
-        auto exprIter = binaryLocations.expressions.find(expr);
-        if (exprIter == binaryLocations.expressions.end()) {
-          // No expression exists for this annotation - perhaps optimizations
-          // removed it.
-          continue;
-        }
-        auto exprOffset = exprIter->second.start;
+        BinaryLocation offset;
+        if (expr == nullptr) {
+          // Function-level annotations have expr==0 and an offset of the start
+          // of the function.
+          offset = 0;
+        } else {
+          auto exprIter = binaryLocations.expressions.find(expr);
+          if (exprIter == binaryLocations.expressions.end()) {
+            // No expression exists for this annotation - perhaps optimizations
+            // removed it.
+            continue;
+          }
+          auto exprOffset = exprIter->second.start;
 
-        if (!funcDeclarationsOffset) {
-          auto funcIter = binaryLocations.functions.find(func.get());
-          assert(funcIter != binaryLocations.functions.end());
-          funcDeclarationsOffset = funcIter->second.declarations;
-        }
+          if (!funcDeclarationsOffset) {
+            auto funcIter = binaryLocations.functions.find(func.get());
+            assert(funcIter != binaryLocations.functions.end());
+            funcDeclarationsOffset = funcIter->second.declarations;
+          }
 
-        // Compute the offset: it should be relative to the start of the
-        // function locals (i.e. the function declarations).
-        auto offset = exprOffset - funcDeclarationsOffset;
+          // Compute the offset: it should be relative to the start of the
+          // function locals (i.e. the function declarations).
+          offset = exprOffset - funcDeclarationsOffset;
+        }
 
         funcHints.exprHints.push_back(ExprHint{expr, offset, &annotation});
       }
@@ -5434,15 +5445,24 @@ void WasmBinaryReader::readExpressionHints(Name sectionName,
 
     auto numHints = getU32LEB();
     for (Index hint = 0; hint < numHints; hint++) {
-      // To get the absolute offset, add the function's offset.
+      // Find the expression this hint is for. If the relative offset is 0, then
+      // it is for the entire function, with expr==null.
+      Expression* expr;
       auto relativeOffset = getU32LEB();
-      auto absoluteOffset = funcLocalsOffset + relativeOffset;
+      if (relativeOffset == 0) {
+        // Function-level annotations have expr==0 and an offset of the start
+        // of the function.
+        expr = nullptr;
+      } else {
+        // To get the absolute offset, add the function's offset.
+        auto absoluteOffset = funcLocalsOffset + relativeOffset;
 
-      auto iter = locationsMap.find(absoluteOffset);
-      if (iter == locationsMap.end()) {
-        throwError("bad offset in " + sectionName.toString());
+        auto iter = locationsMap.find(absoluteOffset);
+        if (iter == locationsMap.end()) {
+          throwError("bad offset in " + sectionName.toString());
+        }
+        expr = iter->second;
       }
-      auto* expr = iter->second;
 
       read(func->codeAnnotations[expr]);
     }

test/lit/inline-hints-func.wast

@@ -0,0 +1,23 @@
+;; RUN: wasm-opt -all             %s -S -o - | filecheck %s
+;; RUN: wasm-opt -all --roundtrip %s -S -o - | filecheck %s
+
+(module
+  (@metadata.code.inline "\12")
+  (func $func-annotation
+    ;; The annotation here is on the function.
+    (drop
+      (i32.const 0)
+    )
+  )
+)
+
+;; CHECK: (module
+;; CHECK-NEXT:  (type $0 (func))
+;; CHECK-NEXT:  (@metadata.code.inline "\12")
+;; CHECK-NEXT:  (func $func-annotation
+;; CHECK-NEXT:   (drop
+;; CHECK-NEXT:    (i32.const 0)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+

test/lit/inline-hints.wast

@@ -82,7 +82,4 @@
       (ref.func $func)
     )
   )
-
-  ;; TODO: test function annotations, after
-  ;; https://github.com/WebAssembly/tool-conventions/issues/251
 )