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

@@ -1506,6 +1506,7 @@ class WasmBinaryWriter {
   };
 
   Module* getModule() { return wasm; }
+  bool getDebugInfo() const { return debugInfo; }
 
   void writeType(Type type);
 
@@ -1569,11 +1570,24 @@ 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;
 
   void prepare();
 
+  // Internal helper for recording the label names of a function once its code
+  // has been written, so that they can be emitted in the name section later.
+  // `labelNames` holds (label index, name) pairs for the function's explicitly
+  // named labels, as gathered by the function writer; its contents are moved
+  // into `funcLabelNames`, and nothing is recorded if it is empty.
+  void noteLabelNames(Function* func,
+                      std::vector<std::pair<Index, Name>>& labelNames);
+
   // Internal helper for emitting a code annotation section for a hint that is
   // expression offset based. Receives the name of the section and two
   // functions, one to check if the annotation we care about exists (receiving
@@ -1721,13 +1735,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 +1760,11 @@ 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;
+
   // before we see a function (like global init expressions), there is no end of
   // function to check
   Index endOfFunction = -1;
@@ -1753,6 +1775,10 @@ class WasmBinaryReader {
   void readFunctions();
   void readVars();
   void setLocalNames(Function& func, Index i);
+  void setLabelNames(Index i);
+  // 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();
 
   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 for `if` and `try_table`,
+  // which cannot hold a name in the IR: a name from the binary name section
+  // should not cause us to introduce a wrapper block that would not otherwise
+  // exist just to hold it.
+  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,18 @@ 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;
+
     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 +680,28 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // its stack.
   std::unordered_map<Name, std::vector<Index>> labelDepths;
 
+  // Wrap the `visitXYZStart` call of a `makeXYZ` function, noting that the
+  // label, if any, was given to us rather than generated by us: only `makeXYZ`
+  // receives labels from outside, that is, from the text format or from the
+  // name section. `visitXYZStart` can also be called directly on an expression
+  // that already exists, as Outlining does when it rebuilds a function, and
+  // the name such an expression carries may be one we generated ourselves.
+  Result<> markLabelExplicit(Name label, Result<> started) {
+    CHECK_ERR(started);
+    if (label && !scopeStack.empty()) {
+      scopeStack.back().labelExplicit = true;
+    }
+    return Ok{};
+  }
+
+  // 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,

src/wasm-stack.h

@@ -129,7 +129,19 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
 
   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;
+
+  // 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());
+
   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 {