[lldb][Expression] Encode Module and DIE UIDs into function AsmLabels (llvm#148877)

LLDB currently attaches `AsmLabel`s to `FunctionDecl`s such that that
the `IRExecutionUnit` can determine which mangled name to call (we can't
rely on Clang deriving the correct mangled name to call because the
debug-info AST doesn't contain all the info that would be encoded in the
DWARF linkage names). However, we don't attach `AsmLabel`s for structors
because they have multiple variants and thus it's not clear which
mangled name to use. In the [RFC on fixing expression evaluation of
abi-tagged
structors](https://discourse.llvm.org/t/rfc-lldb-handling-abi-tagged-constructors-destructors-in-expression-evaluator/82816)
we discussed encoding the structor variant into the `AsmLabel`s.
Specifically in [this
thread](https://discourse.llvm.org/t/rfc-lldb-handling-abi-tagged-constructors-destructors-in-expression-evaluator/82816/7)
we discussed that the contents of the `AsmLabel` are completely under
LLDB's control and we could make use of it to uniquely identify a
function by encoding the exact module and DIE that the function is
associated with (mangled names need not be enough since two identical
mangled symbols may live in different modules). So if we already have a
custom `AsmLabel` format, we can encode the structor variant in a
follow-up (the current idea is to append the structor variant as a
suffix to our custom `AsmLabel` when Clang emits the mangled name into
the JITted IR). Then we would just have to teach the `IRExecutionUnit`
to pick the correct structor variant DIE during symbol resolution. The
draft of this is available
[here](llvm#149827)

This patch sets up the infrastructure for the custom `AsmLabel` format
by encoding the module id, DIE id and mangled name in it.

**Implementation**

The flow is as follows:
1. Create the label in `DWARFASTParserClang`. The format is:
`$__lldb_func:module_id:die_id:mangled_name`
2. When resolving external symbols in `IRExecutionUnit`, we parse this
label and then do a lookup by DIE ID (or mangled name into the module if
the encoded DIE is a declaration).

Depends on llvm#151355

(cherry picked from commit f890591)

Author: Michael137

lldb/include/lldb/Core/Module.h

@@ -88,7 +88,8 @@ struct ModuleFunctionSearchOptions {
 ///
 /// The module will parse more detailed information as more queries are made.
 class Module : public std::enable_shared_from_this<Module>,
-               public SymbolContextScope {
+               public SymbolContextScope,
+               public UserID {
 public:
   class LookupInfo;
   // Static functions that can track the lifetime of module objects. This is

lldb/include/lldb/Core/ModuleList.h

@@ -375,6 +375,14 @@ class ModuleList {
   // UUID values is very efficient and accurate.
   lldb::ModuleSP FindModule(const UUID &uuid) const;
 
+  /// Find a module by LLDB-specific unique identifier.
+  ///
+  /// \param[in] uid The UID of the module assigned to it on construction.
+  ///
+  /// \returns ModuleSP of module with \c uid. Returns nullptr if no such
+  /// module could be found.
+  lldb::ModuleSP FindModule(lldb::user_id_t uid) const;
+
   /// Finds the first module whose file specification matches \a module_spec.
   lldb::ModuleSP FindFirstModule(const ModuleSpec &module_spec) const;
 

lldb/include/lldb/Expression/Expression.h

@@ -13,6 +13,7 @@
 #include <string>
 #include <vector>
 
+#include "llvm/Support/FormatProviders.h"
 
 #include "lldb/Expression/ExpressionTypeSystemHelper.h"
 #include "lldb/lldb-forward.h"
@@ -96,6 +97,62 @@ class Expression {
                                  ///invalid.
 };
 
+/// Holds parsed information about a function call label that
+/// LLDB attaches as an AsmLabel to function AST nodes it parses
+/// from debug-info.
+///
+/// The format being:
+///
+///   <prefix>:<module uid>:<symbol uid>:<name>
+///
+/// The label string needs to stay valid for the entire lifetime
+/// of this object.
+struct FunctionCallLabel {
+  /// Unique identifier of the lldb_private::Module
+  /// which contains the symbol identified by \c symbol_id.
+  lldb::user_id_t module_id;
+
+  /// Unique identifier of the function symbol on which to
+  /// perform the function call. For example, for DWARF this would
+  /// be the DIE UID.
+  lldb::user_id_t symbol_id;
+
+  /// Name to use when searching for the function symbol in
+  /// \c module_id. For most function calls this will be a
+  /// mangled name. In cases where a mangled name can't be used,
+  /// this will be the function name.
+  ///
+  /// NOTE: kept as last element so we don't have to worry about
+  /// ':' in the mangled name when parsing the label.
+  llvm::StringRef lookup_name;
+
+  /// Decodes the specified function \c label into a \c FunctionCallLabel.
+  static llvm::Expected<FunctionCallLabel> fromString(llvm::StringRef label);
+
+  /// Encode this FunctionCallLabel into its string representation.
+  ///
+  /// The representation roundtrips through \c fromString:
+  /// \code{.cpp}
+  /// llvm::StringRef encoded = "$__lldb_func:0x0:0x0:_Z3foov";
+  /// FunctionCallLabel label = *fromString(label);
+  ///
+  /// assert (label.toString() == encoded);
+  /// assert (*fromString(label.toString()) == label);
+  /// \endcode
+  std::string toString() const;
+};
+
+/// LLDB attaches this prefix to mangled names of functions that get called
+/// from JITted expressions.
+inline constexpr llvm::StringRef FunctionCallLabelPrefix = "$__lldb_func";
+
 } // namespace lldb_private
 
+namespace llvm {
+template <> struct format_provider<lldb_private::FunctionCallLabel> {
+  static void format(const lldb_private::FunctionCallLabel &label,
+                     raw_ostream &OS, StringRef Style);
+};
+} // namespace llvm
+
 #endif // LLDB_EXPRESSION_EXPRESSION_H

lldb/include/lldb/Symbol/SymbolFile.h

@@ -349,6 +349,18 @@ class SymbolFile : public PluginInterface {
   GetMangledNamesForFunction(const std::string &scope_qualified_name,
                              std::vector<ConstString> &mangled_names);
 
+  /// Resolves the function corresponding to the specified LLDB function
+  /// call \c label.
+  ///
+  /// \param[in] label The FunctionCallLabel to be resolved.
+  ///
+  /// \returns An llvm::Error if the specified \c label couldn't be resolved.
+  ///          Returns the resolved function (as a SymbolContext) otherwise.
+  virtual llvm::Expected<SymbolContext>
+  ResolveFunctionCallLabel(const FunctionCallLabel &label) {
+    return llvm::createStringError("Not implemented");
+  }
+
   virtual void GetTypes(lldb_private::SymbolContextScope *sc_scope,
                         lldb::TypeClass type_mask,
                         lldb_private::TypeList &type_list) = 0;

lldb/source/Core/Module.cpp

@@ -135,8 +135,10 @@ Module *Module::GetAllocatedModuleAtIndex(size_t idx) {
   return nullptr;
 }
 
+static std::atomic<lldb::user_id_t> g_unique_id = 1;
+
 Module::Module(const ModuleSpec &module_spec)
-    : m_unwind_table(*this), m_file_has_changed(false),
+    : UserID(g_unique_id++), m_unwind_table(*this), m_file_has_changed(false),
       m_first_file_changed_log(false) {
   // Scope for locker below...
   {
@@ -241,7 +243,8 @@ Module::Module(const ModuleSpec &module_spec)
 Module::Module(const FileSpec &file_spec, const ArchSpec &arch,
                ConstString object_name, lldb::offset_t object_offset,
                const llvm::sys::TimePoint<> &object_mod_time)
-    : m_mod_time(FileSystem::Instance().GetModificationTime(file_spec)),
+    : UserID(g_unique_id++),
+      m_mod_time(FileSystem::Instance().GetModificationTime(file_spec)),
       m_arch(arch), m_file(file_spec), m_object_name(object_name),
       m_object_offset(object_offset), m_object_mod_time(object_mod_time),
       m_unwind_table(*this), m_file_has_changed(false),
@@ -262,7 +265,7 @@ Module::Module(const FileSpec &file_spec, const ArchSpec &arch,
 }
 
 Module::Module()
-    : m_unwind_table(*this), m_file_has_changed(false),
+    : UserID(g_unique_id++), m_unwind_table(*this), m_file_has_changed(false),
       m_first_file_changed_log(false) {
   std::lock_guard<std::recursive_mutex> guard(
       GetAllocationModuleCollectionMutex());

lldb/source/Core/ModuleList.cpp

@@ -832,6 +832,20 @@ ModuleSP ModuleList::FindModule(const UUID &uuid) const {
   return module_sp;
 }
 
+ModuleSP ModuleList::FindModule(lldb::user_id_t uid) const {
+  ModuleSP module_sp;
+  ForEach([&](const ModuleSP &m) {
+    if (m->GetID() == uid) {
+      module_sp = m;
+      return IterationAction::Stop;
+    }
+
+    return IterationAction::Continue;
+  });
+
+  return module_sp;
+}
+
 void ModuleList::FindTypes(Module *search_first, const TypeQuery &query,
                            TypeResults &results) const {
   std::lock_guard<std::recursive_mutex> guard(m_modules_mutex);

lldb/source/Expression/Expression.cpp

@@ -10,6 +10,11 @@
 #include "lldb/Target/ExecutionContextScope.h"
 #include "lldb/Target/Target.h"
 
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+
 using namespace lldb_private;
 
 Expression::Expression(Target &target)
@@ -26,3 +31,47 @@ Expression::Expression(ExecutionContextScope &exe_scope)
       m_jit_end_addr(LLDB_INVALID_ADDRESS) {
   assert(m_target_wp.lock());
 }
+
+llvm::Expected<FunctionCallLabel>
+lldb_private::FunctionCallLabel::fromString(llvm::StringRef label) {
+  llvm::SmallVector<llvm::StringRef, 4> components;
+  label.split(components, ":", /*MaxSplit=*/3);
+
+  if (components.size() != 4)
+    return llvm::createStringError("malformed function call label.");
+
+  if (components[0] != FunctionCallLabelPrefix)
+    return llvm::createStringError(llvm::formatv(
+        "expected function call label prefix '{0}' but found '{1}' instead.",
+        FunctionCallLabelPrefix, components[0]));
+
+  llvm::StringRef module_label = components[1];
+  llvm::StringRef die_label = components[2];
+
+  lldb::user_id_t module_id = 0;
+  if (!llvm::to_integer(module_label, module_id))
+    return llvm::createStringError(
+        llvm::formatv("failed to parse module ID from '{0}'.", module_label));
+
+  lldb::user_id_t die_id;
+  if (!llvm::to_integer(die_label, die_id))
+    return llvm::createStringError(
+        llvm::formatv("failed to parse symbol ID from '{0}'.", die_label));
+
+  return FunctionCallLabel{/*.module_id=*/module_id,
+                           /*.symbol_id=*/die_id,
+                           /*.lookup_name=*/components[3]};
+}
+
+std::string lldb_private::FunctionCallLabel::toString() const {
+  return llvm::formatv("{0}:{1:x}:{2:x}:{3}", FunctionCallLabelPrefix,
+                       module_id, symbol_id, lookup_name)
+      .str();
+}
+
+void llvm::format_provider<FunctionCallLabel>::format(
+    const FunctionCallLabel &label, raw_ostream &OS, StringRef Style) {
+  OS << llvm::formatv("FunctionCallLabel{ module_id: {0:x}, symbol_id: {1:x}, "
+                      "lookup_name: {2} }",
+                      label.module_id, label.symbol_id, label.lookup_name);
+}

lldb/source/Expression/IRExecutionUnit.cpp

@@ -15,6 +15,7 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -23,6 +24,7 @@
 #include "lldb/Core/JITSection.h"
 #include "lldb/Core/Module.h"
 #include "lldb/Core/Section.h"
+#include "lldb/Expression/Expression.h"
 #include "lldb/Expression/IRExecutionUnit.h"
 #include "lldb/Expression/ObjectFileJIT.h"
 #include "lldb/Host/HostInfo.h"
@@ -39,6 +41,7 @@
 #include "lldb/Utility/LLDBAssert.h"
 #include "lldb/Utility/LLDBLog.h"
 #include "lldb/Utility/Log.h"
+#include "lldb/lldb-defines.h"
 
 #include <optional>
 
@@ -810,6 +813,40 @@ class LoadAddressResolver {
   lldb::addr_t m_best_internal_load_address = LLDB_INVALID_ADDRESS;
 };
 
+/// Returns address of the function referred to by the special function call
+/// label \c label.
+static llvm::Expected<lldb::addr_t>
+ResolveFunctionCallLabel(const FunctionCallLabel &label,
+                         const lldb_private::SymbolContext &sc,
+                         bool &symbol_was_missing_weak) {
+  symbol_was_missing_weak = false;
+
+  if (!sc.target_sp)
+    return llvm::createStringError("target not available.");
+
+  auto module_sp = sc.target_sp->GetImages().FindModule(label.module_id);
+  if (!module_sp)
+    return llvm::createStringError(
+        llvm::formatv("failed to find module by UID {0}", label.module_id));
+
+  auto *symbol_file = module_sp->GetSymbolFile();
+  if (!symbol_file)
+    return llvm::createStringError(
+        llvm::formatv("no SymbolFile found on module {0:x}.", module_sp.get()));
+
+  auto sc_or_err = symbol_file->ResolveFunctionCallLabel(label);
+  if (!sc_or_err)
+    return llvm::joinErrors(
+        llvm::createStringError("failed to resolve function by UID"),
+        sc_or_err.takeError());
+
+  SymbolContextList sc_list;
+  sc_list.Append(*sc_or_err);
+
+  LoadAddressResolver resolver(*sc.target_sp, symbol_was_missing_weak);
+  return resolver.Resolve(sc_list).value_or(LLDB_INVALID_ADDRESS);
+}
+
 lldb::addr_t
 IRExecutionUnit::FindInSymbols(const std::vector<ConstString> &names,
                                const lldb_private::SymbolContext &sc,
@@ -954,6 +991,34 @@ lldb::addr_t IRExecutionUnit::FindInUserDefinedSymbols(
 
 lldb::addr_t IRExecutionUnit::FindSymbol(lldb_private::ConstString name,
                                          bool &missing_weak) {
+  if (name.GetStringRef().starts_with(FunctionCallLabelPrefix)) {
+    auto label_or_err = FunctionCallLabel::fromString(name);
+    if (!label_or_err) {
+      LLDB_LOG_ERROR(GetLog(LLDBLog::Expressions), label_or_err.takeError(),
+                     "failed to create FunctionCallLabel from '{1}': {0}",
+                     name.GetStringRef());
+      return LLDB_INVALID_ADDRESS;
+    }
+
+    if (auto addr_or_err =
+            ResolveFunctionCallLabel(*label_or_err, m_sym_ctx, missing_weak)) {
+      return *addr_or_err;
+    } else {
+      LLDB_LOG_ERROR(GetLog(LLDBLog::Expressions), addr_or_err.takeError(),
+                     "Failed to resolve function call label '{1}': {0}",
+                     name.GetStringRef());
+
+      // Fall back to lookup by name despite error in resolving the label.
+      // May happen in practice if the definition of a function lives in
+      // a different lldb_private::Module than it's declaration. Meaning
+      // we couldn't pin-point it using the information encoded in the label.
+      name.SetString(label_or_err->lookup_name);
+    }
+  }
+
+  // TODO: now with function call labels, do we still need to
+  // generate alternate manglings?
+
   std::vector<ConstString> candidate_C_names;
   std::vector<ConstString> candidate_CPlusPlus_names;
 

lldb/source/Expression/IRInterpreter.cpp

@@ -259,7 +259,9 @@ class InterpreterStackFrame {
       break;
     case Value::FunctionVal:
       if (const Function *constant_func = dyn_cast<Function>(constant)) {
-        lldb_private::ConstString name(constant_func->getName());
+        lldb_private::ConstString name(
+            llvm::GlobalValue::dropLLVMManglingEscape(
+                constant_func->getName()));
         bool missing_weak = false;
         lldb::addr_t addr = m_execution_unit.FindSymbol(name, missing_weak);
         if (addr == LLDB_INVALID_ADDRESS)

lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.cpp

@@ -26,6 +26,7 @@
 #include "Plugins/Language/ObjC/ObjCLanguage.h"
 #include "lldb/Core/Module.h"
 #include "lldb/Core/Value.h"
+#include "lldb/Expression/Expression.h"
 #include "lldb/Host/Host.h"
 #include "lldb/Symbol/CompileUnit.h"
 #include "lldb/Symbol/CompilerType.h"
@@ -259,7 +260,40 @@ static std::string MakeLLDBFuncAsmLabel(const DWARFDIE &die) {
   if (!name)
     return {};
 
-  return name;
+  SymbolFileDWARF *dwarf = die.GetDWARF();
+  if (!dwarf)
+    return {};
+
+  auto get_module_id = [&](SymbolFile *sym) {
+    if (!sym)
+      return LLDB_INVALID_UID;
+
+    auto *obj = sym->GetMainObjectFile();
+    if (!obj)
+      return LLDB_INVALID_UID;
+
+    auto module_sp = obj->GetModule();
+    if (!module_sp)
+      return LLDB_INVALID_UID;
+
+    return module_sp->GetID();
+  };
+
+  lldb::user_id_t module_id = get_module_id(dwarf->GetDebugMapSymfile());
+  if (module_id == LLDB_INVALID_UID)
+    module_id = get_module_id(dwarf);
+
+  if (module_id == LLDB_INVALID_UID)
+    return {};
+
+  const auto die_id = die.GetID();
+  if (die_id == LLDB_INVALID_UID)
+    return {};
+
+  return FunctionCallLabel{/*module_id=*/module_id,
+                           /*symbol_id=*/die_id,
+                           /*.lookup_name=*/name}
+      .toString();
 }
 
 TypeSP DWARFASTParserClang::ParseTypeFromClangModule(const SymbolContext &sc,