Merge pull request #3051 from augusto2112/parse-reflection-from-jit

[lldb] Implement parsing of metadata generated by the JIT.

Author: augusto2112

lldb/include/lldb/Core/JITSection.h

@@ -0,0 +1,58 @@
+//===-- JITSection.h --------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLDB_CORE_JIT_SECTION_H
+#define LLDB_CORE_JIT_SECTION_H
+
+#include "Section.h"
+
+namespace lldb_private {
+
+/// A JIT-compiled section. This type of section may contain JIT specific data
+/// at the end of its buffer when compared to a regular section.
+class JITSection : public Section {
+public:
+  JITSection(const lldb::ModuleSP &module_sp, ObjectFile *obj_file,
+             lldb::user_id_t sect_id, ConstString name,
+             lldb::SectionType sect_type, lldb::addr_t file_vm_addr,
+             lldb::addr_t vm_size, lldb::offset_t file_offset,
+             lldb::offset_t file_size, uint32_t log2align, uint32_t flags,
+             size_t non_jit_size, uint32_t target_byte_size = 1)
+      : Section(module_sp, obj_file, sect_id, name, sect_type, file_vm_addr,
+                vm_size, file_offset, file_size, log2align, flags,
+                target_byte_size),
+        m_non_jit_size(non_jit_size) {}
+
+  JITSection(const lldb::SectionSP &parent_section_sp,
+             const lldb::ModuleSP &module_sp, ObjectFile *obj_file,
+             lldb::user_id_t sect_id, ConstString name,
+             lldb::SectionType sect_type, lldb::addr_t file_vm_addr,
+             lldb::addr_t vm_size, lldb::offset_t file_offset,
+             lldb::offset_t file_size, uint32_t log2align, uint32_t flags,
+             size_t non_jit_size, uint32_t target_byte_size = 1)
+      : Section(parent_section_sp, module_sp, obj_file, sect_id, name,
+                sect_type, file_vm_addr, vm_size, file_offset, file_size,
+                log2align, flags, target_byte_size),
+        m_non_jit_size(non_jit_size) {}
+
+  // LLVM RTTI support
+  static char ID;
+  virtual bool isA(const void *ClassID) const {
+    return ClassID == &ID || Section::isA(ClassID);
+  }
+  static bool classof(const Section *obj) { return obj->isA(&ID); }
+
+  /// Returns the section size discounting the jit specific data.
+  size_t getNonJitSize() const { return m_non_jit_size; }
+
+private:
+  size_t m_non_jit_size = 0;
+};
+} // namespace lldb_private
+
+#endif // LLDB_CORE_JIT_SECTION_H

lldb/include/lldb/Core/Section.h

@@ -117,7 +117,12 @@ class Section : public std::enable_shared_from_this<Section>,
           lldb::offset_t file_size, uint32_t log2align, uint32_t flags,
           uint32_t target_byte_size = 1);
 
-  ~Section();
+  virtual ~Section();
+
+  // LLVM RTTI support
+  static char ID;
+  virtual bool isA(const void *ClassID) const { return ClassID == &ID; }
+  static bool classof(const Section *obj) { return obj->isA(&ID); }
 
   static int Compare(const Section &a, const Section &b);
 

lldb/include/lldb/Expression/IRExecutionUnit.h

@@ -17,6 +17,7 @@
 
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Module.h"
+#include "llvm/ADT/StringMap.h"
 
 #include "lldb/Expression/IRMemoryMap.h"
 #include "lldb/Symbol/ObjectFile.h"
@@ -409,6 +410,8 @@ class IRExecutionUnit : public std::enable_shared_from_this<IRExecutionUnit>,
   ///< defining no functions using that variable, would do this.)  If this
   ///< is true, any allocations need to be committed immediately -- no
   ///< opportunity for relocation.
+
+  llvm::StringMap<uint64_t> m_section_size_map;
 };
 
 } // namespace lldb_private

lldb/source/Core/CMakeLists.txt

@@ -39,6 +39,7 @@ add_lldb_library(lldbCore
   Highlighter.cpp
   IOHandler.cpp
   IOHandlerCursesGUI.cpp
+  JITSection.cpp
   Mangled.cpp
   Module.cpp
   ModuleChild.cpp

lldb/source/Core/JITSection.cpp

@@ -0,0 +1,14 @@
+//===-- JITSection.cpp ----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+
+#include "lldb/Core/JITSection.h"
+
+namespace lldb_private {
+char JITSection::ID;
+}

lldb/source/Core/Section.cpp

@@ -25,6 +25,8 @@ class DataExtractor;
 using namespace lldb;
 using namespace lldb_private;
 
+char Section::ID;
+
 const char *Section::GetTypeAsCString() const {
   switch (m_type) {
   case eSectionTypeInvalid:

lldb/source/Expression/IRExecutionUnit.cpp

@@ -6,7 +6,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
 #include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
@@ -16,6 +18,7 @@
 
 #include "lldb/Core/Debugger.h"
 #include "lldb/Core/Disassembler.h"
+#include "lldb/Core/JITSection.h"
 #include "lldb/Core/Module.h"
 #include "lldb/Core/Section.h"
 #include "lldb/Expression/IRExecutionUnit.h"
@@ -336,6 +339,30 @@ void IRExecutionUnit::GetRunnableInfo(Status &error, lldb::addr_t &func_addr,
     m_execution_engine_up->setObjectCache(m_object_cache_up.get());
   }
 
+  /// Record the size of the sections discounting the extra space used 
+  /// by the JIT. This is a hack and should be removed when we switch
+  /// from MCJIT to ORCJIT.
+  struct JitSectionSizeRecorder : public llvm::JITEventListener {
+    llvm::StringMap<uint64_t> &m_section_size_map;
+    JitSectionSizeRecorder(llvm::StringMap<uint64_t> &section_size_map) : 
+      m_section_size_map(section_size_map) {}
+
+    void notifyObjectLoaded(ObjectKey K, const llvm::object::ObjectFile &Obj,
+                            const llvm::RuntimeDyld::LoadedObjectInfo &L) override {
+      for (llvm::object::SectionRef section: Obj.sections()) {
+        auto name = section.getName();
+        if (name)
+          m_section_size_map.insert({name.get(), section.getSize()});
+      }
+      
+    }
+  };
+
+  JitSectionSizeRecorder size_recorder(m_section_size_map);
+  m_execution_engine_up->RegisterJITEventListener(&size_recorder);
+  auto on_exit = llvm::make_scope_exit(
+      [&]() { m_execution_engine_up->UnregisterJITEventListener(&size_recorder); });
+
   // Make sure we see all sections, including ones that don't have
   // relocations...
   m_execution_engine_up->setProcessAllSections(true);
@@ -1325,15 +1352,17 @@ void IRExecutionUnit::PopulateSectionList(
     lldb_private::SectionList &section_list) {
   for (AllocationRecord &record : m_records) {
     if (record.m_size > 0) {
-      lldb::SectionSP section_sp(new lldb_private::Section(
+      size_t non_jit_size = m_section_size_map.lookup(record.m_name);
+      lldb::SectionSP section_sp(new lldb_private::JITSection(
           obj_file->GetModule(), obj_file, record.m_section_id,
           ConstString(record.m_name), record.m_sect_type,
           record.m_process_address, record.m_size,
           record.m_host_address, // file_offset (which is the host address for
                                  // the data)
           record.m_size,         // file_size
           0,
-          record.m_permissions)); // flags
+          record.m_permissions, // flags
+          non_jit_size)); 
       section_list.AddSection(section_sp);
     }
   }

lldb/source/Plugins/ExpressionParser/Swift/SwiftUserExpression.cpp

@@ -352,6 +352,8 @@ bool SwiftUserExpression::Parse(DiagnosticManager &diagnostic_manager,
   const lldb::LanguageType lang_type = lldb::eLanguageTypeSwift;
 
   m_options.SetLanguage(lang_type);
+  m_options.SetGenerateDebugInfo(generate_debug_info);
+  
   uint32_t first_body_line = 0;
 
   // I have to pass some value for add_locals.  I'm passing "false" because
@@ -473,7 +475,14 @@ bool SwiftUserExpression::Parse(DiagnosticManager &diagnostic_manager,
     StreamString jit_module_name;
     jit_module_name.Printf("%s%u", FunctionName(),
                            m_options.GetExpressionNumber());
-    m_execution_unit_sp->CreateJITModule(jit_module_name.GetString().data());
+    auto module = m_execution_unit_sp->CreateJITModule(jit_module_name.GetString().data());
+
+    auto *swift_runtime = SwiftLanguageRuntime::Get(process);
+    if (module && swift_runtime) {
+      ModuleList modules;
+      modules.Append(module, false);
+      swift_runtime->ModulesDidLoad(modules);
+    }
   }
 
   if (jit_error.Success()) {

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.cpp

@@ -18,6 +18,7 @@
 #include "Utility/ARM64_DWARF_Registers.h"
 #include "lldb/Breakpoint/StoppointCallbackContext.h"
 #include "lldb/Core/Debugger.h"
+#include "lldb/Core/JITSection.h"
 #include "lldb/Core/PluginManager.h"
 #include "lldb/Core/Section.h"
 #include "lldb/Core/ValueObjectConstResult.h"
@@ -543,6 +544,58 @@ void SwiftLanguageRuntime::ModulesDidLoad(const ModuleList &module_list) {
   }
 }
 
+bool SwiftLanguageRuntimeImpl::AddJitObjectFileToReflectionContext(
+    ObjectFile &obj_file, llvm::Triple::ObjectFormatType obj_format_type) {
+  assert(obj_file.GetType() == ObjectFile::eTypeJIT &&
+         "Not a JIT object file!");
+  std::unique_ptr<swift::SwiftObjectFileFormat> obj_file_format;
+  switch (obj_format_type) {
+  case llvm::Triple::MachO:
+    obj_file_format = std::make_unique<swift::SwiftObjectFileFormatMachO>();
+    break;
+  case llvm::Triple::ELF:
+    obj_file_format = std::make_unique<swift::SwiftObjectFileFormatELF>();
+    break;
+  case llvm::Triple::COFF:
+    obj_file_format = std::make_unique<swift::SwiftObjectFileFormatCOFF>();
+    break;
+  default:
+    if (Log *log = GetLogIfAnyCategoriesSet(LIBLLDB_LOG_TYPES))
+      log->Printf("%s: Could not find out swift reflection section names for "
+                  "object format type.",
+                  __FUNCTION__);
+    return false;
+  }
+
+  return m_reflection_ctx->addImage(
+      [&](swift::ReflectionSectionKind section_kind)
+          -> std::pair<swift::remote::RemoteRef<void>, uint64_t> {
+        auto section_name =
+            ConstString(obj_file_format->getSectionName(section_kind));
+        for (auto section : *obj_file.GetSectionList()) {
+          JITSection *jit_section = llvm::dyn_cast<JITSection>(section.get());
+          if (jit_section && section->GetName() == section_name) {
+            DataExtractor extractor;
+            auto section_size = section->GetSectionData(extractor);
+            if (!section_size)
+              return {};
+            auto size = jit_section->getNonJitSize();
+            auto data = extractor.GetData();
+            if (section_size < size || !data.begin())
+              return {};
+
+            auto *Buf = malloc(size);
+            std::memcpy(Buf, data.begin(), size);
+            swift::remote::RemoteRef<void> remote_ref(section->GetFileAddress(),
+                                                      Buf);
+
+            return {remote_ref, size};
+          }
+        }
+        return {};
+      });
+}
+
 bool SwiftLanguageRuntimeImpl::AddModuleToReflectionContext(
     const lldb::ModuleSP &module_sp) {
   // This function is called from within SetupReflection so it cannot
@@ -558,6 +611,12 @@ bool SwiftLanguageRuntimeImpl::AddModuleToReflectionContext(
   Address start_address = obj_file->GetBaseAddress();
   auto load_ptr = static_cast<uintptr_t>(
       start_address.GetLoadAddress(&(m_process.GetTarget())));
+  if (obj_file->GetType() == ObjectFile::eTypeJIT) {
+    auto object_format_type =
+        module_sp->GetArchitecture().GetTriple().getObjectFormat();
+    return AddJitObjectFileToReflectionContext(*obj_file, object_format_type);
+  }
+
   if (load_ptr == 0 || load_ptr == LLDB_INVALID_ADDRESS) {
     if (obj_file->GetType() != ObjectFile::eTypeJIT)
       if (Log *log = GetLogIfAnyCategoriesSet(LIBLLDB_LOG_TYPES))

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntimeDynamicTypeResolution.cpp

@@ -1805,73 +1805,6 @@ bool SwiftLanguageRuntimeImpl::IsValidErrorValue(ValueObject &in_value) {
 
   return true;
 }
-/// Digs into the protocols inside a protocol composition and checks if 
-/// any of them belong to the LLDB module.
-static bool
-IsDeclaredInLLDBExpr(const swift::reflection::TypeRef *typeref) {
-  using swift::Demangle::Node;
-
-  const auto *protocol_typeref = llvm::dyn_cast_or_null<
-      const swift::reflection::ProtocolCompositionTypeRef>(typeref);
-
-  if (!protocol_typeref)
-    return false;
-
-  swift::Demangle::Demangler dem;
-
-  NodePointer type = protocol_typeref->getDemangling(dem);
-  if (!type || type->getKind() != Node::Kind::Type ||
-      type->getNumChildren() == 0)
-    return false;
-
-  NodePointer maybe_protocol_list = type->getFirstChild();
-  NodePointer protocol_list;
-
-  if (!maybe_protocol_list)
-    return false;
-
-  switch (maybe_protocol_list->getKind()) {
-  case Node::Kind::ProtocolList:
-    protocol_list = maybe_protocol_list;
-    break;
-  case Node::Kind::ProtocolListWithAnyObject:
-  case Node::Kind::ProtocolListWithClass:
-    if (maybe_protocol_list->getNumChildren() == 0)
-      return false;
-    protocol_list = maybe_protocol_list->getFirstChild();
-    break;
-  default:
-    return false;
-  }
-
-  if (!protocol_list || protocol_list->getKind() != Node::Kind::ProtocolList ||
-      protocol_list->getNumChildren() == 0)
-    return false;
-
-  NodePointer type_list = protocol_list->getFirstChild();
-  if (!type_list || type_list->getKind() != Node::Kind::TypeList)
-    return false;
-
-  for (NodePointer internal_type : *type_list) {
-    if (!internal_type || internal_type->getKind() != Node::Kind::Type ||
-        internal_type->getNumChildren() == 0)
-      continue;
-
-    NodePointer protocol = internal_type->getFirstChild();
-    if (!protocol || protocol->getKind() != Node::Kind::Protocol ||
-        protocol->getNumChildren() == 0)
-      continue;
-
-    NodePointer module = protocol->getFirstChild();
-    if (!module || module->getKind() != Node::Kind::Module ||
-        !module->hasText())
-      continue;
-
-    if (module->getText().startswith("__lldb_expr_"))
-      return true;
-  }
-  return false;
-}
 
 bool SwiftLanguageRuntimeImpl::GetDynamicTypeAndAddress_Protocol(
     ValueObject &in_value, CompilerType protocol_type,
@@ -1957,17 +1890,6 @@ bool SwiftLanguageRuntimeImpl::GetDynamicTypeAndAddress_Protocol(
     return false;
   }
 
-  // If the protocol list contains any types declared in LLDB,
-  // we use the remote AST implementation since we currently
-  // don't register the metadata generated by LLDB.
-  if (IsDeclaredInLLDBExpr(protocol_typeref)) {
-    auto pair = remote_ast_impl(use_local_buffer, existential_address);
-    if (!pair)
-      return false;
-    class_type_or_name.SetCompilerType(std::get<CompilerType>(*pair));
-    address = std::get<Address>(*pair);
-    return true;
-  }
 
   if (use_local_buffer)
     PushLocalBuffer(existential_address, in_value.GetByteSize().getValueOr(0));