[lldb][Expression] Encode Module and DIE UIDs into function AsmLabels

Author: Michael137

lldb/include/lldb/Core/Module.h

@@ -86,7 +86,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
@@ -97,6 +98,7 @@ class Module : public std::enable_shared_from_this<Module>,
   // using the "--global" (-g for short).
   static size_t GetNumberAllocatedModules();
 
+  static Module *GetAllocatedModuleWithUID(lldb::user_id_t uid);
   static Module *GetAllocatedModuleAtIndex(size_t idx);
 
   static std::recursive_mutex &GetAllocationModuleCollectionMutex();

lldb/include/lldb/Expression/Expression.h

@@ -96,6 +96,41 @@ 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>:<mangled name>:<module id>:<DIE id>
+///
+/// The label string needs to stay valid for the entire lifetime
+/// of this object.
+struct FunctionCallLabel {
+  /// Mostly for debuggability.
+  llvm::StringRef m_pubname;
+
+  lldb::user_id_t m_die_id;
+  lldb::user_id_t m_module_id;
+};
+
+/// LLDB attaches this prefix to mangled names of functions that it get called
+/// from JITted expressions.
+inline constexpr llvm::StringRef FunctionCallLabelPrefix = "$__lldb_func";
+
+bool consumeFunctionCallLabelPrefix(llvm::StringRef &name);
+bool hasFunctionCallLabelPrefix(llvm::StringRef name);
+
+/// Returns the components of the specified function call label.
+///
+/// The format of \c label is described in \c FunctionCallLabel.
+/// The label prefix is not one of the components.
+llvm::Expected<llvm::SmallVector<llvm::StringRef, 3>>
+splitFunctionCallLabel(llvm::StringRef label);
+
+// Decodes the function label into a \c FunctionCallLabel.
+llvm::Expected<FunctionCallLabel> makeFunctionCallLabel(llvm::StringRef label);
+
 } // namespace lldb_private
 
 #endif // LLDB_EXPRESSION_EXPRESSION_H

lldb/include/lldb/Symbol/SymbolFile.h

@@ -18,6 +18,7 @@
 #include "lldb/Symbol/CompilerType.h"
 #include "lldb/Symbol/Function.h"
 #include "lldb/Symbol/SourceModule.h"
+#include "lldb/Symbol/SymbolContext.h"
 #include "lldb/Symbol/Type.h"
 #include "lldb/Symbol/TypeList.h"
 #include "lldb/Symbol/TypeSystem.h"
@@ -328,6 +329,19 @@ class SymbolFile : public PluginInterface {
   GetMangledNamesForFunction(const std::string &scope_qualified_name,
                              std::vector<ConstString> &mangled_names);
 
+  /// Resolves the function DIE uniquely identified by \c uid within
+  /// this SymbolFile.
+  ///
+  /// \param[in,out] sc_list The resolved functions will be appended to this
+  /// list.
+  ///
+  /// \param[in] uid The UID of the function DIE to resolve.
+  ///
+  virtual llvm::Error ResolveFunctionUID(SymbolContextList &sc_list,
+                                         lldb::user_id_t uid) {
+    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

@@ -121,6 +121,15 @@ size_t Module::GetNumberAllocatedModules() {
   return GetModuleCollection().size();
 }
 
+Module *Module::GetAllocatedModuleWithUID(lldb::user_id_t uid) {
+  std::lock_guard<std::recursive_mutex> guard(
+      GetAllocationModuleCollectionMutex());
+  for (Module *mod : GetModuleCollection())
+    if (mod->GetID() == uid)
+      return mod;
+  return nullptr;
+}
+
 Module *Module::GetAllocatedModuleAtIndex(size_t idx) {
   std::lock_guard<std::recursive_mutex> guard(
       GetAllocationModuleCollectionMutex());
@@ -130,8 +139,11 @@ Module *Module::GetAllocatedModuleAtIndex(size_t idx) {
   return nullptr;
 }
 
+// TODO: needs a mutex
+static 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...
   {
@@ -236,7 +248,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),
@@ -257,7 +270,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/Expression/Expression.cpp

@@ -10,6 +10,10 @@
 #include "lldb/Target/ExecutionContextScope.h"
 #include "lldb/Target/Target.h"
 
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+
 using namespace lldb_private;
 
 Expression::Expression(Target &target)
@@ -26,3 +30,64 @@ Expression::Expression(ExecutionContextScope &exe_scope)
       m_jit_end_addr(LLDB_INVALID_ADDRESS) {
   assert(m_target_wp.lock());
 }
+
+bool lldb_private::consumeFunctionCallLabelPrefix(llvm::StringRef &name) {
+  // On Darwin mangled names get a '_' prefix.
+  name.consume_front("_");
+  return name.consume_front(FunctionCallLabelPrefix);
+}
+
+bool lldb_private::hasFunctionCallLabelPrefix(llvm::StringRef name) {
+  // On Darwin mangled names get a '_' prefix.
+  name.consume_front("_");
+  return name.starts_with(FunctionCallLabelPrefix);
+}
+
+// Expected format is:
+// $__lldb_func:<mangled name>:<module id>:<definition/declaration DIE id>
+llvm::Expected<llvm::SmallVector<llvm::StringRef, 3>>
+lldb_private::splitFunctionCallLabel(llvm::StringRef label) {
+  if (!consumeFunctionCallLabelPrefix(label))
+    return llvm::createStringError(
+        "expected function call label prefix not found in %s", label.data());
+
+  if (!label.consume_front(":"))
+    return llvm::createStringError(
+        "incorrect format: expected ':' as the first character.");
+
+  llvm::SmallVector<llvm::StringRef, 3> components;
+  label.split(components, ":");
+
+  if (components.size() != 3)
+    return llvm::createStringError(
+        "incorrect format: too many label subcomponents.");
+
+  return components;
+}
+
+llvm::Expected<FunctionCallLabel>
+lldb_private::makeFunctionCallLabel(llvm::StringRef label) {
+  auto components_or_err = splitFunctionCallLabel(label);
+  if (!components_or_err)
+    return llvm::joinErrors(
+        llvm::createStringError("Failed to decode function call label"),
+        components_or_err.takeError());
+
+  const auto &components = *components_or_err;
+
+  llvm::StringRef module_label = components[1];
+  llvm::StringRef die_label = components[2];
+
+  lldb::user_id_t module_id = 0;
+  if (module_label.consumeInteger(0, module_id))
+    return llvm::createStringError(llvm::formatv(
+        "failed to parse module ID from '%s'.", components[1].data()));
+
+  lldb::user_id_t die_id;
+  if (die_label.consumeInteger(/*Radix=*/0, die_id))
+    return llvm::createStringError("failed to parse DIE ID from '%s'.",
+                                   components[2].data());
+
+  return FunctionCallLabel{
+      .m_pubname = components[0], .m_die_id = die_id, .m_module_id = module_id};
+}

lldb/source/Expression/IRExecutionUnit.cpp

@@ -13,13 +13,15 @@
 #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"
 
 #include "lldb/Core/Debugger.h"
 #include "lldb/Core/Disassembler.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"
@@ -36,6 +38,7 @@
 #include "lldb/Utility/LLDBAssert.h"
 #include "lldb/Utility/LLDBLog.h"
 #include "lldb/Utility/Log.h"
+#include "lldb/lldb-defines.h"
 
 #include <optional>
 
@@ -771,6 +774,52 @@ 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.
+///
+/// \param[in] label Function call label encoding the unique location of the
+/// function to look up.
+///                  Assumes that the \c FunctionCallLabelPrefix has been
+///                  stripped from the front of the label.
+static llvm::Expected<lldb::addr_t>
+ResolveFunctionCallLabel(llvm::StringRef name,
+                         const lldb_private::SymbolContext &sc,
+                         bool &symbol_was_missing_weak) {
+  symbol_was_missing_weak = false;
+
+  auto label_or_err = makeFunctionCallLabel(name);
+  if (!label_or_err)
+    return llvm::joinErrors(
+        llvm::createStringError("failed to create FunctionCallLabel from: %s",
+                                name.data()),
+        label_or_err.takeError());
+
+  const auto &label = *label_or_err;
+
+  Module *module = Module::GetAllocatedModuleWithUID(label.m_module_id);
+
+  if (!module)
+    return llvm::createStringError(
+        llvm::formatv("failed to find module by UID {0}", label.m_module_id));
+
+  auto *symbol_file = module->GetSymbolFile();
+  if (!symbol_file)
+    return llvm::createStringError(
+        llvm::formatv("no SymbolFile found on module {0:x}.", module));
+
+  SymbolContextList sc_list;
+  if (auto err = symbol_file->ResolveFunctionUID(sc_list, label.m_die_id))
+    return llvm::joinErrors(
+        llvm::createStringError("failed to resolve function by UID"),
+        std::move(err));
+
+  if (!sc.target_sp)
+    return llvm::createStringError("target not available.");
+
+  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,
@@ -906,6 +955,20 @@ lldb::addr_t IRExecutionUnit::FindInUserDefinedSymbols(
 
 lldb::addr_t IRExecutionUnit::FindSymbol(lldb_private::ConstString name,
                                          bool &missing_weak) {
+  if (hasFunctionCallLabelPrefix(name.GetStringRef())) {
+    if (auto addr_or_err = ResolveFunctionCallLabel(name.GetStringRef(),
+                                                    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());
+      return LLDB_INVALID_ADDRESS;
+    }
+  }
+
+  // TODO: do we still need the following lookups?
+
   std::vector<ConstString> candidate_C_names;
   std::vector<ConstString> candidate_CPlusPlus_names;
 

lldb/source/Plugins/ExpressionParser/Clang/ClangExpressionDeclMap.cpp

@@ -1991,7 +1991,7 @@ void ClangExpressionDeclMap::AddContextClassType(NameSearchContext &context,
     const bool is_artificial = false;
 
     CXXMethodDecl *method_decl = m_clang_ast_context->AddMethodToCXXRecordType(
-        copied_clang_type.GetOpaqueQualType(), "$__lldb_expr", nullptr,
+        copied_clang_type.GetOpaqueQualType(), "$__lldb_expr", std::nullopt,
         method_type, lldb::eAccessPublic, is_virtual, is_static, is_inline,
         is_explicit, is_attr_used, is_artificial);