Read and write label names from the name section

CHANGELOG.md

@@ -17,6 +17,7 @@ Current Trunk
 
  - Add a new wasm-embed tool for extracting Wasm modules embedded as byte arrays
    in JS source files or replacing such modules in the JS source (#9127)
+ - Support reading and writing label names in the name section (#9166)
 
 v133
 ----

src/binaryen-c.cpp

@@ -5537,6 +5537,27 @@ void BinaryenStringSliceWTFSetEnd(BinaryenExpressionRef expr,
 
 // Functions
 
+// Every label name in a body built through the C API was given to us by the
+// caller, so all of them are explicit.
+static void noteExplicitLabelNames(Function* func) {
+  if (!func->body) {
+    return;
+  }
+  struct Scanner : public PostWalker<Scanner> {
+    std::unordered_set<Name>& names;
+    Scanner(std::unordered_set<Name>& names) : names(names) {}
+    void note(Name name) {
+      if (name) {
+        names.insert(name);
+      }
+    }
+    void visitBlock(Block* curr) { note(curr->name); }
+    void visitLoop(Loop* curr) { note(curr->name); }
+    void visitTry(Try* curr) { note(curr->name); }
+  } scanner(func->explicitLabelNames);
+  scanner.walk(func->body);
+}
+
 static BinaryenFunctionRef addFunctionInternal(BinaryenModuleRef module,
                                                const char* name,
                                                HeapType type,
@@ -5550,6 +5571,7 @@ static BinaryenFunctionRef addFunctionInternal(BinaryenModuleRef module,
     ret->vars.push_back(Type(varTypes[i]));
   }
   ret->body = (Expression*)body;
+  noteExplicitLabelNames(ret);
 
   // Lock. This can be called from multiple threads at once, and is a
   // point where they all access and modify the module.
@@ -6594,6 +6616,7 @@ void BinaryenFunctionSetBody(BinaryenFunctionRef func,
                              BinaryenExpressionRef body) {
   assert(body);
   ((Function*)func)->body = (Expression*)body;
+  noteExplicitLabelNames((Function*)func);
 }
 BinaryenHeapType BinaryenFunctionGetType(BinaryenFunctionRef func) {
   return ((Function*)func)->type.getHeapType().getID();

src/ir/module-utils.cpp

@@ -74,6 +74,7 @@ copyFunctionWithoutAdd(Function* func,
   ret->vars = func->vars;
   ret->localNames = func->localNames;
   ret->localIndices = func->localIndices;
+  ret->explicitLabelNames = func->explicitLabelNames;
   ret->body = ExpressionManipulator::copy(func->body, out);
   metadata::copyBetweenFunctions(func->body, ret->body, func, ret.get());
   ret->funcAnnotations = func->funcAnnotations;

src/wasm-binary.h

@@ -1438,6 +1438,8 @@ class WasmBinaryWriter {
   void writeFunctionSignatures();
   void writeExpression(Expression* curr);
   void writeFunctions();
+  void noteLabelNames(Function* func,
+                      std::vector<std::pair<Index, Name>>& labelNames);
   void writeStrings();
   void writeGlobals();
   void writeExports();
@@ -1506,6 +1508,7 @@ class WasmBinaryWriter {
   };
 
   Module* getModule() { return wasm; }
+  bool getDebugInfo() const { return debugInfo; }
 
   void writeType(Type type);
 
@@ -1569,6 +1572,11 @@ class WasmBinaryWriter {
   // info here, and then use it when writing the names.
   std::unordered_map<Name, MappedLocals> funcMappedLocals;
 
+  // The explicitly named labels of each function, as (label index, name),
+  // gathered while writing the code section and used to write the name
+  // section afterwards.
+  std::unordered_map<Name, std::vector<std::pair<Index, Name>>> funcLabelNames;
+
   // Indexes in the string literal section of each StringConst in the wasm.
   std::unordered_map<Name, Index> stringIndexes;
 
@@ -1721,13 +1729,16 @@ class WasmBinaryReader {
   Signature getSignatureByTypeIndex(Index index);
   Signature getSignatureByFunctionIndex(Index index);
 
-  Name getNextLabel();
-
   // We read the names section first so we know in advance what names various
   // elements should have. Store the information for use when building
   // expressions.
   std::unordered_map<Index, Name> functionNames;
   std::unordered_map<Index, std::unordered_map<Index, Name>> localNames;
+  // Label names, indexed by function index and then by the index of the label
+  // in the function. Labels are indexed in the order the instructions that
+  // introduce them appear in the function body, including the ones that have
+  // no name in the name section.
+  std::unordered_map<Index, std::unordered_map<Index, Name>> labelNames;
   std::unordered_map<Index, Name> typeNames;
   std::unordered_map<Index, std::unordered_map<Index, Name>> fieldNames;
   std::unordered_map<Index, Name> tableNames;
@@ -1743,6 +1754,14 @@ class WasmBinaryReader {
     usedGlobalNames, usedTagNames;
 
   Function* currFunction = nullptr;
+  // The label names of the function we are currently reading, if it has any,
+  // and the index of the next label in it.
+  const std::unordered_map<Index, Name>* currLabelNames = nullptr;
+  Index nextLabelIndex = 0;
+
+  // Returns the name the name section gives to the next label of the current
+  // function, or a null name if it has none, and advances the label index.
+  Name getNextLabelName();
   // before we see a function (like global init expressions), there is no end of
   // function to check
   Index endOfFunction = -1;
@@ -1753,6 +1772,7 @@ class WasmBinaryReader {
   void readFunctions();
   void readVars();
   void setLocalNames(Function& func, Index i);
+  void setLabelNames(Index i);
 
   Result<> readInst();
 

src/wasm-ir-builder.h

@@ -105,6 +105,13 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
                               Type inputType = Type::none);
   Result<> visitEnd();
 
+  // Set the name to use for the label of the scope we are currently building,
+  // if that scope ends up needing a label. This is used for names that come
+  // from the binary name section, which are optional: unlike names parsed from
+  // the text format, they should not cause us to introduce blocks that would
+  // not otherwise exist just to have something to hold the name.
+  void setScopeNameHint(Name name);
+
   // Used to visit break nodes when traversing a single block without its
   // context. The type indicates how many values the break carries to its
   // destination.
@@ -414,8 +421,25 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     // for branches, since the normal label is only used for delegates.
     Name branchLabel;
 
+    // The name to use if we end up needing a label for this scope. Unlike
+    // `label`, this does not itself cause the scope to be labeled. See
+    // `setScopeNameHint`.
+    Name nameHint;
+
+    // For Try scopes with a label in the text format, the fresh version of that
+    // label, reserved for whichever of `label` and `branchLabel` is needed
+    // first. Most tries are only targeted by branches, and giving the name to
+    // the delegate label up front would leave the branch label with a
+    // generated name.
+    Name reservedLabel;
+
     bool labelUsed = false;
 
+    // Whether this scope's label was given explicitly (in the text format or in
+    // the binary name section) rather than generated by us. Recorded in the
+    // function's `explicitLabelNames`.
+    bool labelExplicit = false;
+
     // If the control flow scope has an input type, we need to lower it using a
     // scratch local because we cannot represent control flow input in the IR.
     Type inputType;
@@ -663,6 +687,14 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // its stack.
   std::unordered_map<Name, std::vector<Index>> labelDepths;
 
+  // Note that a label was given to us rather than generated by us, so that we
+  // know it is worth writing out to the name section.
+  void noteExplicitLabel(Name label) {
+    if (func && label) {
+      func->explicitLabelNames.insert(label);
+    }
+  }
+
   Name makeFresh(Name label, Index hint = 0) {
     return Names::getValidName(
       label,
@@ -678,10 +710,16 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
 
   Result<> pushScope(ScopeCtx&& scope) {
     if (auto label = scope.getOriginalLabel()) {
-      // Assign a fresh label to the scope, if necessary.
-      if (!scope.label) {
+      // Assign a fresh label to the scope, if necessary. The labels of a try
+      // are only assigned when they are used, and catches keep the labels of
+      // their try.
+      if (scope.getTry()) {
+        scope.reservedLabel = makeFresh(label);
+        scope.nameHint = label;
+      } else if (!scope.label && !scope.getCatch() && !scope.getCatchAll()) {
         scope.label = makeFresh(label);
       }
+      scope.labelExplicit = true;
       // Record the original label to handle references to it correctly.
       labelDepths[label].push_back(scopeStack.size() + 1);
     }

src/wasm-stack.h

@@ -127,9 +127,21 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
   void emitUnreachableLocalSet(Index index);
   void mapLocalsAndEmitHeader();
 
+  // Note the start of a scope: every block, loop, if, try and try_table takes
+  // an index in the function's label index space, in the order they appear in
+  // the function body, whether or not they have a name. Records the explicitly
+  // named ones for the name section.
+  void noteScopeStart(Name name = Name());
+
   MappedLocals mappedLocals;
 
+  // The explicitly named labels of this function, as (label index, name).
+  std::vector<std::pair<Index, Name>> labelNames;
+
 private:
+  // The index the next scope in this function will take.
+  Index nextLabelIndex = 0;
+
   void emitMemoryAccess(size_t alignment,
                         size_t bytes,
                         uint64_t offset,
@@ -565,6 +577,9 @@ class BinaryenIRToBinaryWriter
   }
 
   MappedLocals& getMappedLocals() { return writer.mappedLocals; }
+  std::vector<std::pair<Index, Name>>& getLabelNames() {
+    return writer.labelNames;
+  }
 
 private:
   WasmBinaryWriter& parent;
@@ -607,6 +622,9 @@ class StackIRToBinaryWriter {
   void write();
 
   MappedLocals& getMappedLocals() { return writer.mappedLocals; }
+  std::vector<std::pair<Index, Name>>& getLabelNames() {
+    return writer.labelNames;
+  }
 
 private:
   WasmBinaryWriter& parent;

src/wasm.h

@@ -33,6 +33,7 @@
 #include <ostream>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "ir/import-names.h"
@@ -2448,6 +2449,15 @@ class Function : public Importable {
   std::unordered_map<Index, Name> localNames;
   std::unordered_map<Name, Index> localIndices;
 
+  // The label names in this function that were given explicitly - in the text
+  // format, in the binary name section, or through the C API - as opposed to
+  // the ones we generated ourselves. Only these are worth writing out to the
+  // name section. Label names are unique inside a function, so this keeps
+  // identifying the right labels even as optimizations replace the expressions
+  // that carry them. A name in here may no longer exist in the body, which is
+  // harmless.
+  std::unordered_set<Name> explicitLabelNames;
+
   // Source maps debugging info: map expression nodes to their file, line, col,
   // symbol name.
   struct DebugLocation {