[lld][WebAssembly] Allow --trace-symbol to work with symbols with custom import names

Author: sbc100

lld/test/wasm/trace-symbol.s

@@ -1,19 +1,28 @@
+# Test -y symbol and -trace-symbol=symbol
+
 # RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %p/Inputs/ret32.s -o %t.ret32.o
 # RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.start.o %s
 # RUN: wasm-ld -o %t.wasm %t.start.o %t.ret32.o -y ret32 -y _start | FileCheck %s -check-prefix=BOTH
 # RUN: wasm-ld -o %t.wasm %t.ret32.o %t.start.o -y ret32 -y _start | FileCheck %s -check-prefix=REVERSED
 
-# check alias
+# check long argument form
 # RUN: wasm-ld -o %t.wasm %t.start.o %t.ret32.o -trace-symbol=_start | FileCheck %s -check-prefixes=JUST-START
 
+# check import names can be traced
+# RUN: wasm-ld -o %t.wasm %t.start.o %t.ret32.o -trace-symbol=foo_external | FileCheck %s -check-prefixes=IMPORT-NAME
+
 .functype ret32 (f32) -> (i32)
 
+.functype foo_import () -> ()
+.import_name foo_import, foo_external
+
 .globl  _start
 _start:
   .functype _start () -> ()
   f32.const 0.0
   call ret32
   drop
+  call foo_import
   end_function
 
 # BOTH:          start.o: definition of _start
@@ -26,3 +35,6 @@ _start:
 
 # JUST-START: start.o: definition of _start
 # JUST-START-NOT: ret32
+
+# IMPORT-NAME: start.o: reference to foo_external
+# IMPORT-NAME-NOT: ret32

lld/wasm/SymbolTable.cpp

@@ -507,6 +507,15 @@ static void setImportAttributes(T *existing,
   }
 }
 
+static void traceImport(std::optional<StringRef> importName, InputFile *file) {
+  if (importName.has_value()) {
+    auto name = importName.value();
+    if (symtab->isTraced(name)) {
+      printTraceSymbolUndefined(name, file);
+    }
+  }
+}
+
 Symbol *SymbolTable::addUndefinedFunction(StringRef name,
                                           std::optional<StringRef> importName,
                                           std::optional<StringRef> importModule,
@@ -526,6 +535,7 @@ Symbol *SymbolTable::addUndefinedFunction(StringRef name,
     printTraceSymbolUndefined(name, file);
 
   auto replaceSym = [&]() {
+    traceImport(importName, file);
     replaceSymbol<UndefinedFunction>(s, name, importName, importModule, flags,
                                      file, sig, isCalledDirectly);
   };
@@ -560,6 +570,7 @@ Symbol *SymbolTable::addUndefinedFunction(StringRef name,
         replaceSym();
     }
     if (existingUndefined) {
+      traceImport(importName, file);
       setImportAttributes(existingUndefined, importName, importModule, flags,
                           file);
       if (isCalledDirectly)
@@ -612,10 +623,11 @@ Symbol *SymbolTable::addUndefinedGlobal(StringRef name,
   if (s->traced)
     printTraceSymbolUndefined(name, file);
 
-  if (wasInserted)
+  if (wasInserted) {
+    traceImport(importName, file);
     replaceSymbol<UndefinedGlobal>(s, name, importName, importModule, flags,
                                    file, type);
-  else if (auto *lazy = dyn_cast<LazySymbol>(s))
+  } else if (auto *lazy = dyn_cast<LazySymbol>(s))
     lazy->extract();
   else if (s->isDefined())
     checkGlobalType(s, file, type);
@@ -638,10 +650,11 @@ Symbol *SymbolTable::addUndefinedTable(StringRef name,
   if (s->traced)
     printTraceSymbolUndefined(name, file);
 
-  if (wasInserted)
+  if (wasInserted) {
+    traceImport(importName, file);
     replaceSymbol<UndefinedTable>(s, name, importName, importModule, flags,
                                   file, type);
-  else if (auto *lazy = dyn_cast<LazySymbol>(s))
+  } else if (auto *lazy = dyn_cast<LazySymbol>(s))
     lazy->extract();
   else if (s->isDefined())
     checkTableType(s, file, type);
@@ -664,10 +677,11 @@ Symbol *SymbolTable::addUndefinedTag(StringRef name,
   if (s->traced)
     printTraceSymbolUndefined(name, file);
 
-  if (wasInserted)
+  if (wasInserted) {
+    traceImport(importName, file);
     replaceSymbol<UndefinedTag>(s, name, importName, importModule, flags, file,
                                 sig);
-  else if (auto *lazy = dyn_cast<LazySymbol>(s))
+  } else if (auto *lazy = dyn_cast<LazySymbol>(s))
     lazy->extract();
   else if (s->isDefined())
     checkTagType(s, file, sig);
@@ -828,9 +842,20 @@ bool SymbolTable::getFunctionVariant(Symbol* sym, const WasmSignature *sig,
 // Set a flag for --trace-symbol so that we can print out a log message
 // if a new symbol with the same name is inserted into the symbol table.
 void SymbolTable::trace(StringRef name) {
+  tracingEnabled = true;
   symMap.insert({CachedHashStringRef(name), -1});
 }
 
+bool SymbolTable::isTraced(StringRef name) {
+  // Early exit in the default case to avoid symbol hashmap lookup.
+  if (!tracingEnabled)
+    return false;
+  auto it = symMap.find(CachedHashStringRef(name));
+  if (it == symMap.end())
+    return false;
+  return it->second == -1 || symVector[it->second]->traced;
+}
+
 void SymbolTable::wrap(Symbol *sym, Symbol *real, Symbol *wrap) {
   // Swap symbols as instructed by -wrap.
   int &origIdx = symMap[CachedHashStringRef(sym->getName())];

lld/wasm/SymbolTable.h

@@ -50,6 +50,8 @@ class SymbolTable {
 
   void trace(StringRef name);
 
+  bool isTraced(StringRef name);
+
   Symbol *addDefinedFunction(StringRef name, uint32_t flags, InputFile *file,
                              InputFunction *function);
   Symbol *addDefinedData(StringRef name, uint32_t flags, InputFile *file,
@@ -132,6 +134,8 @@ class SymbolTable {
 
   // For LTO.
   std::unique_ptr<BitcodeCompiler> lto;
+
+  bool tracingEnabled = false;
 };
 
 extern SymbolTable *symtab;