Implement parsing of metadata generated by the JIT.

TypeSystemSwiftTyperef requires metadata to be registered for it to be
able to do its work. This registers metadata generated by the JIT, which
are types declared in LLDB itself (including the REPL and playgrounds).

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));