[NFC] Refactor function-level annotations to a separate field

src/ir/intrinsics.h

@@ -127,7 +127,7 @@ class Intrinsics {
 
   // Get the code annotations for a function itself.
   static CodeAnnotation getAnnotations(Function* func) {
-    return getAnnotations(nullptr, func);
+    return func->funcAnnotations;
   }
 
   void setAnnotations(Function* func,
@@ -137,7 +137,7 @@ class Intrinsics {
   }
 
   void setAnnotations(Function* func, const CodeAnnotation& value) {
-    setAnnotations(func, nullptr, value);
+    func->funcAnnotations = value;
   }
 
   // Given a call in a function, return all the annotations for it. The call may

src/ir/metadata.cpp

@@ -76,10 +76,7 @@ void copyBetweenFunctions(Expression* origin,
   }
 
   // Also copy function-level annotations, if any.
-  auto iter = originAnnotations.find(nullptr);
-  if (iter != originAnnotations.end()) {
-    copyAnnotations[nullptr] = iter->second;
-  }
+  copyFunc->funcAnnotations = originFunc->funcAnnotations;
 }
 
 #pragma GCC diagnostic push

src/parser/contexts.h

@@ -1482,10 +1482,8 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
         Builder::addVar(f.get(), l.name, l.type);
       }
     }
-    // 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);
+      f->funcAnnotations = parseAnnotations(annotations);
     }
     return Ok{};
   }

src/passes/Print.cpp

@@ -279,6 +279,7 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
   // Print code annotations for an expression. If the expression is nullptr,
   // prints for the current function.
   void printCodeAnnotations(Expression* curr);
+  void printCodeAnnotations(const CodeAnnotation& annotation);
 
   void printExpressionContents(Expression* curr);
 
@@ -2772,43 +2773,46 @@ void PrintSExpression::printDebugDelimiterLocation(Expression* curr, Index i) {
 void PrintSExpression::printCodeAnnotations(Expression* curr) {
   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);
-    }
-    if (annotation.inline_) {
-      Colors::grey(o);
-      std::ofstream saved;
-      saved.copyfmt(o);
-      o << "(@" << Annotations::InlineHint << " \"\\" << std::hex
-        << std::setfill('0') << std::setw(2) << int(*annotation.inline_)
-        << "\")\n";
-      o.copyfmt(saved);
-      restoreNormalColor(o);
-      doIndent(o, indent);
-    }
-    if (annotation.removableIfUnused) {
-      Colors::grey(o);
-      o << "(@" << Annotations::RemovableIfUnusedHint << ")\n";
-      restoreNormalColor(o);
-      doIndent(o, indent);
-    }
-    if (annotation.jsCalled) {
-      Colors::grey(o);
-      o << "(@" << Annotations::JSCalledHint << ")\n";
-      restoreNormalColor(o);
-      doIndent(o, indent);
-    }
-    if (annotation.idempotent) {
-      Colors::grey(o);
-      o << "(@" << Annotations::IdempotentHint << ")\n";
-      restoreNormalColor(o);
-      doIndent(o, indent);
-    }
+    printCodeAnnotations(iter->second);
+  }
+}
+
+void PrintSExpression::printCodeAnnotations(const CodeAnnotation& annotation) {
+  if (annotation.branchLikely) {
+    Colors::grey(o);
+    o << "(@" << Annotations::BranchHint << " \"\\0"
+      << (*annotation.branchLikely ? "1" : "0") << "\")\n";
+    restoreNormalColor(o);
+    doIndent(o, indent);
+  }
+  if (annotation.inline_) {
+    Colors::grey(o);
+    std::ofstream saved;
+    saved.copyfmt(o);
+    o << "(@" << Annotations::InlineHint << " \"\\" << std::hex
+      << std::setfill('0') << std::setw(2) << int(*annotation.inline_)
+      << "\")\n";
+    o.copyfmt(saved);
+    restoreNormalColor(o);
+    doIndent(o, indent);
+  }
+  if (annotation.removableIfUnused) {
+    Colors::grey(o);
+    o << "(@" << Annotations::RemovableIfUnusedHint << ")\n";
+    restoreNormalColor(o);
+    doIndent(o, indent);
+  }
+  if (annotation.jsCalled) {
+    Colors::grey(o);
+    o << "(@" << Annotations::JSCalledHint << ")\n";
+    restoreNormalColor(o);
+    doIndent(o, indent);
+  }
+  if (annotation.idempotent) {
+    Colors::grey(o);
+    o << "(@" << Annotations::IdempotentHint << ")\n";
+    restoreNormalColor(o);
+    doIndent(o, indent);
   }
 }
 
@@ -3266,7 +3270,7 @@ void PrintSExpression::visitImportedFunction(Function* curr) {
   lastPrintedLocation = std::nullopt;
   o << '(';
   emitImportHeader(curr);
-  printCodeAnnotations(nullptr);
+  printCodeAnnotations(curr->funcAnnotations);
   handleSignature(curr);
   o << "))";
   o << maybeNewLine;
@@ -3280,7 +3284,7 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
   if (currFunction->prologLocation) {
     printDebugLocation(*currFunction->prologLocation);
   }
-  printCodeAnnotations(nullptr);
+  printCodeAnnotations(curr->funcAnnotations);
   handleSignature(curr, true);
   incIndent();
   for (size_t i = curr->getVarIndexBase(); i < curr->getNumLocals(); i++) {

src/passes/StripToolchainAnnotations.cpp

@@ -37,14 +37,13 @@ struct StripToolchainAnnotations
   }
 
   void doWalkFunction(Function* func) {
+    remove(func->funcAnnotations);
+
     auto& annotations = func->codeAnnotations;
     auto iter = annotations.begin();
     while (iter != annotations.end()) {
-      // Remove the toolchain-specific annotations.
       auto& annotation = iter->second;
-      annotation.removableIfUnused = false;
-      annotation.jsCalled = false;
-      annotation.idempotent = false;
+      remove(annotation);
 
       // If nothing remains, remove the entire annotation.
       if (annotation == CodeAnnotation()) {
@@ -54,6 +53,13 @@ struct StripToolchainAnnotations
       }
     }
   }
+
+  // Remove all toolchain-specific annotations.
+  void remove(CodeAnnotation& annotation) {
+    annotation.removableIfUnused = false;
+    annotation.jsCalled = false;
+    annotation.idempotent = false;
+  }
 };
 
 Pass* createStripToolchainAnnotationsPass() {

src/wasm.h

@@ -2334,18 +2334,17 @@ class Function : public Importable {
     delimiterLocations;
   BinaryLocations::FunctionLocations funcLocation;
 
-  // 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
+  // Annotations on expressions. 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.
-  // XXX As an unordered map, if this is modified by one thread, another should
-  //     not be reading it. That should not happen atm - all annotations are
-  //     set up in dedicated passes or in the binary reader - but if one pass
-  //     could add an expression annotation, another should not at the same time
-  //     read the function-level annotations, even though that is natural to do.
-  //     We may want to move the function-level annotations to a dedicated
-  //     field outside the map.
   std::unordered_map<Expression*, CodeAnnotation> codeAnnotations;
+  // Annotations on the function itself. These could be stored in the above map
+  // with a key of nullptr (matching the binary format), but it is safer to
+  // keep them separate: in theory a parallel pass could modify code annotations
+  // by, say, duplicating code and adding new ones, which modifies the map,
+  // while another thread might query that function, for whom it has a call,
+  // about the function-level annotations.
+  CodeAnnotation funcAnnotations;
 
   // The effects for this function, if they have been computed. We use a shared
   // ptr here to avoid compilation errors with the forward-declared

src/wasm/wasm-binary.cpp

@@ -1661,35 +1661,33 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeExpressionHints(
 
     for (auto& [expr, annotation] : func->codeAnnotations) {
       if (has(annotation)) {
-        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;
-          }
+        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;
 
-          // Compute the offset: it should be relative to the start of the
-          // function locals (i.e. the function declarations).
-          offset = exprOffset - funcDeclarationsOffset;
+        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;
+
         funcHints.exprHints.push_back(ExprHint{expr, offset, &annotation});
       }
     }
 
+    auto& funcAnnotations = func->funcAnnotations;
+    if (has(funcAnnotations)) {
+      funcHints.exprHints.push_back(ExprHint{nullptr, 0, &funcAnnotations});
+    }
+
     if (funcHints.exprHints.empty()) {
       continue;
     }
@@ -5495,23 +5493,21 @@ void WasmBinaryReader::readExpressionHints(Name sectionName,
     for (Index hint = 0; hint < numHints; hint++) {
       // 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();
       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());
-        }
-        expr = iter->second;
+        // Function-level annotations have the offset of the start of the
+        // function.
+        read(func->funcAnnotations);
+        continue;
       }
+      // 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* expr = iter->second;
       read(func->codeAnnotations[expr]);
     }
   }