[lldb/Swift] Fix import attributes handing in following expression evaluation

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