[lldb/Swift] Fix import attributes handing in expression evaluations

Previously, when importing a Swift module it the REPL or a Playground, LLDB
would only cache the module name and reload the module at every expression
evaluation. Doing so means that, when LLDB transfers imports information from
the original source file to the new one, it is doing it with low fidelity,
because LLDB would not keep track of the import attributes used in the previous
expression evaluations.

Thanks to swiftlang/swift#33935, we can now cache the whole
`swift::AttributedImport<swift::ImportedModule>>` instead of only storing the
module name. This struct contains both the `ImportedModule` with the `ModuleDecl`
but also the `ImportOptions` set containing all the possible import attributes
(Exported / Testable / PrivateImport ...) of that import statement.

Now, before compiling and evaluating a Swift expression, LLDB fetches all the
attributed imports cached previously and add them to the `ImplicitImportInfo`
object that is used to create the `ModuleDecl` of the evaluated expression.
It allows for the `ImportOptions` to be propagated in the compiled expression.

This should fix some irregularities between the debugger and the
compiler regarding import statements.

rdar://65319954

Signed-off-by: Med Ismail Bennani <[email protected]>

Author: medismailben

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

@@ -55,17 +55,18 @@
 #include "clang/Rewrite/Core/RewriteBuffer.h"
 
 #include "swift/AST/ASTContext.h"
-#include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticConsumer.h"
+#include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/IRGenOptions.h"
 #include "swift/AST/IRGenRequests.h"
+#include "swift/AST/Import.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ModuleLoader.h"
-#include "swift/Demangling/Demangle.h"
+#include "swift/Basic/OptimizationMode.h"
 #include "swift/Basic/PrimarySpecificPaths.h"
 #include "swift/Basic/SourceManager.h"
-#include "swift/Basic/OptimizationMode.h"
 #include "swift/ClangImporter/ClangImporter.h"
+#include "swift/Demangling/Demangle.h"
 #include "swift/Frontend/Frontend.h"
 #include "swift/Parse/LocalContext.h"
 #include "swift/Parse/PersistentParserState.h"
@@ -1137,6 +1138,7 @@ static llvm::Expected<ParsedExpression> ParseAndImport(
     lldb::StackFrameWP &stack_frame_wp, SymbolContext &sc,
     ExecutionContextScope &exe_scope, const EvaluateExpressionOptions &options,
     bool repl, bool playground) {
+  Log *log(lldb_private::GetLogIfAllCategoriesSet(LIBLLDB_LOG_EXPRESSIONS));
 
   auto should_disable_objc_runtime = [&]() {
     lldb::StackFrameSP this_frame_sp(stack_frame_wp.lock());
@@ -1162,7 +1164,20 @@ static llvm::Expected<ParsedExpression> ParseAndImport(
   if (playground) {
     auto *persistent_state =
         sc.target_sp->GetSwiftPersistentExpressionState(exe_scope);
-    persistent_state->AddHandLoadedModule(ConstString("Swift"));
+
+    Status error;
+    SourceModule module_info;
+    module_info.path.emplace_back("Swift");
+    swift::ModuleDecl *module =
+        swift_ast_context->GetModule(module_info, error);
+
+    if (error.Fail() || !module) {
+      LLDB_LOG(log, "couldn't load Swift Standard Library\n");
+      return error.ToError();
+    }
+
+    persistent_state->AddHandLoadedModule(ConstString("Swift"),
+                                          swift::ImportedModule(module));
   }
 
   std::string main_filename;
@@ -1178,7 +1193,8 @@ static llvm::Expected<ParsedExpression> ParseAndImport(
   // The Swift stdlib needs to be imported before the SwiftLanguageRuntime can
   // be used.
   Status implicit_import_error;
-  llvm::SmallVector<swift::ModuleDecl *, 16> additional_imports;
+  llvm::SmallVector<swift::AttributedImport<swift::ImportedModule>, 16>
+      additional_imports;
   if (!SwiftASTContext::GetImplicitImports(*swift_ast_context, sc, exe_scope,
                                            stack_frame_wp, additional_imports,
                                            implicit_import_error)) {
@@ -1188,9 +1204,8 @@ static llvm::Expected<ParsedExpression> ParseAndImport(
 
   swift::ImplicitImportInfo importInfo;
   importInfo.StdlibKind = swift::ImplicitStdlibKind::Stdlib;
-  for (auto *module : additional_imports)
-    importInfo.AdditionalImports.emplace_back(swift::ImportedModule(module),
-                                              swift::ImportOptions());
+  for (auto &attributed_import : additional_imports)
+    importInfo.AdditionalImports.emplace_back(attributed_import);
 
   auto module_id = ast_context->getIdentifier(expr_name_buf);
   auto &module = *swift::ModuleDecl::create(module_id, *ast_context,

lldb/source/Plugins/ExpressionParser/Swift/SwiftPersistentExpressionState.h

@@ -15,10 +15,13 @@
 
 #include "SwiftExpressionVariable.h"
 
+#include "swift/AST/Import.h"
+#include "swift/AST/Module.h"
+
 #include "lldb/Core/SwiftForward.h"
 #include "lldb/Expression/ExpressionVariable.h"
 
-#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 
 #include <set>
@@ -38,7 +41,9 @@ namespace lldb_private {
 /// 0-based counter for naming result variables.
 //----------------------------------------------------------------------
 class SwiftPersistentExpressionState : public PersistentExpressionState {
-  typedef std::set<lldb_private::ConstString> HandLoadedModuleSet;
+
+  typedef llvm::StringMap<swift::AttributedImport<swift::ImportedModule>>
+      HandLoadedModuleSet;
 
 public:
   class SwiftDeclMap {
@@ -114,8 +119,11 @@ class SwiftPersistentExpressionState : public PersistentExpressionState {
 
   // This just adds this module to the list of hand-loaded modules, it doesn't
   // actually load it.
-  void AddHandLoadedModule(ConstString module_name) {
-    m_hand_loaded_modules.insert(module_name);
+  void AddHandLoadedModule(
+      ConstString module_name,
+      swift::AttributedImport<swift::ImportedModule> attributed_import) {
+    m_hand_loaded_modules.insert_or_assign(module_name.GetStringRef(),
+                                           attributed_import);
   }
 
   /// This returns the list of hand-loaded modules.

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

@@ -285,8 +285,38 @@ bool SwiftUserExpression::Parse(DiagnosticManager &diagnostic_manager,
                                  "couldn't start parsing (no target)");
     return false;
   }
+
+  StackFrame *frame = exe_ctx.GetFramePtr();
+  if (!frame) {
+    LLDB_LOG(log, "no stack frame");
+    return false;
+  }
+
+  // Make sure the target's SwiftASTContext has been setup before doing any
+  // Swift name lookups.
+  llvm::Optional<SwiftASTContextReader> maybe_swift_ast_ctx =
+      target->GetScratchSwiftASTContext(err, *frame);
+  if (!maybe_swift_ast_ctx) {
+    LLDB_LOG(log, "no Swift AST Context");
+    return false;
+  }
+
+  SwiftASTContext *swift_ast_ctx = maybe_swift_ast_ctx->get();
+
   if (auto *persistent_state = GetPersistentState(target, exe_ctx)) {
-    persistent_state->AddHandLoadedModule(ConstString("Swift"));
+
+    Status error;
+    SourceModule module_info;
+    module_info.path.emplace_back("Swift");
+    swift::ModuleDecl *module = swift_ast_ctx->GetModule(module_info, error);
+
+    if (error.Fail() || !module) {
+      LLDB_LOG(log, "couldn't load Swift Standard Library\n");
+      return false;
+    }
+
+    persistent_state->AddHandLoadedModule(ConstString("Swift"),
+                                          swift::ImportedModule(module));
     m_result_delegate.RegisterPersistentState(persistent_state);
     m_error_delegate.RegisterPersistentState(persistent_state);
   } else {

lldb/source/Plugins/TypeSystem/Swift/SwiftASTContext.cpp

@@ -8284,7 +8284,9 @@ static swift::ModuleDecl *LoadOneModule(const SourceModule &module,
 bool SwiftASTContext::GetImplicitImports(
     SwiftASTContext &swift_ast_context, SymbolContext &sc,
     ExecutionContextScope &exe_scope, lldb::StackFrameWP &stack_frame_wp,
-    llvm::SmallVectorImpl<swift::ModuleDecl *> &modules, Status &error) {
+    llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+        &modules,
+    Status &error) {
   if (!GetCompileUnitImports(swift_ast_context, sc, stack_frame_wp, modules,
                              error)) {
     return false;
@@ -8294,15 +8296,28 @@ bool SwiftASTContext::GetImplicitImports(
       sc.target_sp->GetSwiftPersistentExpressionState(exe_scope);
 
   // Get the hand-loaded modules from the SwiftPersistentExpressionState.
-  for (ConstString name : persistent_expression_state->GetHandLoadedModules()) {
+  for (auto &module_pair :
+       persistent_expression_state->GetHandLoadedModules()) {
+
+    auto &attributed_import = module_pair.second;
+
+    // If the ImportedModule in the SwiftPersistentExpressionState has a
+    // non-null ModuleDecl, add it to the ImplicitImports list.
+    if (attributed_import.module.importedModule) {
+      modules.emplace_back(attributed_import);
+      continue;
+    }
+
+    // Otherwise, try reloading the ModuleDecl using the module name.
     SourceModule module_info;
-    module_info.path.push_back(name);
-    auto *module = LoadOneModule(module_info, swift_ast_context, stack_frame_wp,
-                                 error);
+    module_info.path.emplace_back(module_pair.first());
+    auto *module =
+        LoadOneModule(module_info, swift_ast_context, stack_frame_wp, error);
     if (!module)
       return false;
 
-    modules.push_back(module);
+    attributed_import.module = swift::ImportedModule(module);
+    modules.emplace_back(attributed_import);
   }
   return true;
 }
@@ -8314,7 +8329,6 @@ bool SwiftASTContext::CacheUserImports(SwiftASTContext &swift_ast_context,
                                        swift::SourceFile &source_file,
                                        Status &error) {
   llvm::SmallString<1> m_description;
-  llvm::SmallVector<swift::ImportedModule, 2> parsed_imports;
 
   swift::ModuleDecl::ImportFilter import_filter {
       swift::ModuleDecl::ImportFilterKind::Exported,
@@ -8324,13 +8338,12 @@ bool SwiftASTContext::CacheUserImports(SwiftASTContext &swift_ast_context,
       swift::ModuleDecl::ImportFilterKind::ShadowedByCrossImportOverlay
   };
 
-  source_file.getImportedModules(parsed_imports, import_filter);
-
   auto *persistent_expression_state =
       sc.target_sp->GetSwiftPersistentExpressionState(exe_scope);
 
-  for (auto module_pair : parsed_imports) {
-    swift::ModuleDecl *module = module_pair.importedModule;
+  for (const auto &attributed_import : source_file.getImports()) {
+    swift::ModuleDecl *module = attributed_import.module.importedModule;
+
     if (module) {
       std::string module_name;
       GetNameFromModule(module, module_name);
@@ -8346,7 +8359,8 @@ bool SwiftASTContext::CacheUserImports(SwiftASTContext &swift_ast_context,
           return false;
 
         // How do we tell we are in REPL or playground mode?
-        persistent_expression_state->AddHandLoadedModule(module_const_str);
+        persistent_expression_state->AddHandLoadedModule(module_const_str,
+                                                         attributed_import);
       }
     }
   }
@@ -8356,16 +8370,18 @@ bool SwiftASTContext::CacheUserImports(SwiftASTContext &swift_ast_context,
 bool SwiftASTContext::GetCompileUnitImports(
     SwiftASTContext &swift_ast_context, SymbolContext &sc,
     lldb::StackFrameWP &stack_frame_wp,
-    llvm::SmallVectorImpl<swift::ModuleDecl *> &modules, Status &error) {
+    llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+        &modules,
+    Status &error) {
   // Import the Swift standard library and its dependencies.
   SourceModule swift_module;
-  swift_module.path.push_back(ConstString("Swift"));
+  swift_module.path.emplace_back("Swift");
   auto *stdlib =
       LoadOneModule(swift_module, swift_ast_context, stack_frame_wp, error);
   if (!stdlib)
     return false;
 
-  modules.push_back(stdlib);
+  modules.emplace_back(swift::ImportedModule(stdlib));
 
   CompileUnit *compile_unit = sc.comp_unit;
   if (!compile_unit || compile_unit->GetLanguage() != lldb::eLanguageTypeSwift)
@@ -8386,7 +8402,7 @@ bool SwiftASTContext::GetCompileUnitImports(
     if (!loaded_module)
       return false;
 
-    modules.push_back(loaded_module);
+    modules.emplace_back(swift::ImportedModule(loaded_module));
   }
   return true;
 }

lldb/source/Plugins/TypeSystem/Swift/SwiftASTContext.h

@@ -750,7 +750,9 @@ class SwiftASTContext : public TypeSystemSwift {
   static bool GetImplicitImports(
       SwiftASTContext &swift_ast_context, SymbolContext &sc,
       ExecutionContextScope &exe_scope, lldb::StackFrameWP &stack_frame_wp,
-      llvm::SmallVectorImpl<swift::ModuleDecl *> &modules, Status &error);
+      llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+          &modules,
+      Status &error);
 
   /// Cache the user's imports from a SourceFile in a given execution scope such
   /// that they are carried over into future expression evaluations.
@@ -761,11 +763,12 @@ class SwiftASTContext : public TypeSystemSwift {
                                swift::SourceFile &source_file, Status &error);
 
   /// Retrieve the modules imported by the compilation unit.
-  static bool
-  GetCompileUnitImports(SwiftASTContext &swift_ast_context, SymbolContext &sc,
-                        lldb::StackFrameWP &stack_frame_wp,
-                        llvm::SmallVectorImpl<swift::ModuleDecl *> &modules,
-                        Status &error);
+  static bool GetCompileUnitImports(
+      SwiftASTContext &swift_ast_context, SymbolContext &sc,
+      lldb::StackFrameWP &stack_frame_wp,
+      llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+          &modules,
+      Status &error);
 
 protected:
   /// This map uses the string value of ConstStrings as the key, and the

lldb/source/Target/Target.cpp

@@ -2504,7 +2504,8 @@ llvm::Optional<SwiftASTContextReader> Target::GetScratchSwiftASTContext(
       !swift_ast_ctx->HasFatalErrors()) {
     StackFrameWP frame_wp(frame_sp);
     SymbolContext sc = frame_sp->GetSymbolContext(lldb::eSymbolContextEverything);
-    llvm::SmallVector<swift::ModuleDecl *, 16> modules;
+    llvm::SmallVector<swift::AttributedImport<swift::ImportedModule>, 16>
+        modules;
     swift_ast_ctx->GetCompileUnitImports(*swift_ast_ctx, sc, frame_wp, modules,
                                          error);
   }

lldb/test/Shell/SwiftREPL/Inputs/Test.swift

@@ -0,0 +1,11 @@
+public struct Bar {
+    public init(baz: Int) { self.baz = baz }
+
+    var baz : Int
+}
+
+struct Foo {
+    init(bar: Int) { self.bar = bar }
+
+    var bar : Int
+}

lldb/test/Shell/SwiftREPL/LookupWithAttributedImport.test

@@ -0,0 +1,24 @@
+// Make sure import attributes are taken into consideration at every evaluation
+// rdar://65319954
+// REQUIRES: swift
+
+// RUN: rm -rf %t
+// RUN: mkdir %t
+// RUN: %target-swiftc -emit-library %S/Inputs/Test.swift -module-name Test -emit-module-path %t/Test.swiftmodule -o %t/libTest%target-shared-library-suffix -Onone -g  -enable-testing
+
+// RUN: %lldb --repl="-I%t -L%t -lTest" < %s 2>&1 | FileCheck %s
+
+import Test
+
+let y = Bar(baz: 123)
+// CHECK: y: Test.Bar = {
+// CHECK: baz = 123
+
+let x = Foo(bar:42)
+// CHECK: error: repl.swift:{{.*}}: error: cannot find 'Foo' in scope
+
+@testable import Test
+
+let x = Foo(bar:42)
+// CHECK: x: Test.Foo = {
+// CHECK: bar = 42