Don't eagerly import all impls. (#4447)

Instead, eagerly import only the impls from the api file corresponding
to the current file, if any, because we need those for impl
redeclaration lookup. For all other cases, load only the impls in
libraries that are referenced as part of an impl lookup query.

toolchain/check/check.cpp

@@ -923,9 +923,9 @@ static auto CheckParseTree(
 
   InitPackageScopeAndImports(context, unit_info, total_ir_count);
 
-  // Import all impls declared in imports.
-  // TODO: Do this selectively when we see an impl query.
-  ImportImpls(context);
+  // Eagerly import the impls declared in the api file to prepare to redeclare
+  // them.
+  ImportImplsFromApiFile(context);
 
   if (!ProcessNodeIds(context, vlog_stream, unit_info.err_tracker,
                       node_converters[unit_info.check_ir_id.index])) {

toolchain/check/impl_lookup.cpp

@@ -7,13 +7,103 @@
 #include "toolchain/check/deduce.h"
 #include "toolchain/check/generic.h"
 #include "toolchain/check/import_ref.h"
+#include "toolchain/sem_ir/ids.h"
 
 namespace Carbon::Check {
 
+static auto FindAssociatedImportIRs(Context& context,
+                                    SemIR::ConstantId type_const_id,
+                                    SemIR::ConstantId interface_const_id)
+    -> llvm::SmallVector<SemIR::ImportIRId> {
+  llvm::SmallVector<SemIR::ImportIRId> result;
+
+  // Add an entity to our result.
+  auto add_entity = [&](const SemIR::EntityWithParamsBase& entity) {
+    // We will look for impls in the import IR associated with the first owning
+    // declaration.
+    auto decl_id = entity.first_owning_decl_id;
+    if (!decl_id.is_valid()) {
+      return;
+    }
+    auto loc_id = context.insts().GetLocId(decl_id);
+    if (loc_id.is_import_ir_inst_id()) {
+      result.push_back(
+          context.import_ir_insts().Get(loc_id.import_ir_inst_id()).ir_id);
+    }
+  };
+
+  llvm::SmallVector<SemIR::InstId> worklist;
+  worklist.push_back(context.constant_values().GetInstId(type_const_id));
+  worklist.push_back(context.constant_values().GetInstId(interface_const_id));
+
+  while (!worklist.empty()) {
+    auto inst_id = worklist.pop_back_val();
+
+    // Visit the operands of the constant.
+    auto inst = context.insts().Get(inst_id);
+    auto [arg0_kind, arg1_kind] = inst.ArgKinds();
+    for (auto [arg, kind] :
+         {std::pair{inst.arg0(), arg0_kind}, {inst.arg1(), arg1_kind}}) {
+      switch (kind) {
+        case SemIR::IdKind::For<SemIR::InstId>: {
+          if (auto id = SemIR::InstId(arg); id.is_valid()) {
+            worklist.push_back(id);
+          }
+          break;
+        }
+        case SemIR::IdKind::For<SemIR::InstBlockId>: {
+          if (auto id = SemIR::InstBlockId(arg); id.is_valid()) {
+            auto block = context.inst_blocks().Get(id);
+            worklist.append(block.begin(), block.end());
+          }
+          break;
+        }
+        case SemIR::IdKind::For<SemIR::ClassId>: {
+          add_entity(context.classes().Get(SemIR::ClassId(arg)));
+          break;
+        }
+        case SemIR::IdKind::For<SemIR::InterfaceId>: {
+          add_entity(context.interfaces().Get(SemIR::InterfaceId(arg)));
+          break;
+        }
+        case SemIR::IdKind::For<SemIR::FunctionId>: {
+          add_entity(context.functions().Get(SemIR::FunctionId(arg)));
+          break;
+        }
+        default: {
+          break;
+        }
+      }
+    }
+  }
+
+  // Deduplicate.
+  std::sort(result.begin(), result.end(),
+            [](SemIR::ImportIRId a, SemIR::ImportIRId b) {
+              return a.index < b.index;
+            });
+  result.erase(std::unique(result.begin(), result.end()), result.end());
+
+  return result;
+}
+
 auto LookupInterfaceWitness(Context& context, SemIR::LocId loc_id,
                             SemIR::ConstantId type_const_id,
                             SemIR::ConstantId interface_const_id)
     -> SemIR::InstId {
+  auto import_irs =
+      FindAssociatedImportIRs(context, type_const_id, interface_const_id);
+  for (auto import_ir : import_irs) {
+    // TODO: Instead of importing all impls, only import ones that are in some
+    // way connected to this query.
+    for (auto impl_index : llvm::seq(
+             context.import_irs().Get(import_ir).sem_ir->impls().size())) {
+      // TODO: Track the relevant impls and only consider those ones and any
+      // local impls, rather than looping over all impls below.
+      ImportImpl(context, import_ir, SemIR::ImplId(impl_index));
+    }
+  }
+
   // TODO: Add a better impl lookup system. At the very least, we should only be
   // considering impls that are for the same interface we're querying. We can
   // also skip impls that mention any types that aren't part of our impl query.

toolchain/check/import_ref.cpp

@@ -2382,24 +2382,29 @@ auto LoadImportRef(Context& context, SemIR::InstId inst_id) -> void {
   }
 }
 
-// TODO: This doesn't belong in this file. Consider moving the import resolver
-// and this file elsewhere.
-auto ImportImpls(Context& context) -> void {
-  for (auto [import_index, import_ir] :
-       llvm::enumerate(context.import_irs().array_ref())) {
-    if (!import_ir.sem_ir) {
-      continue;
-    }
+auto ImportImplsFromApiFile(Context& context) -> void {
+  SemIR::ImportIRId import_ir_id = SemIR::ImportIRId::ApiForImpl;
+  auto& import_ir = context.import_irs().Get(import_ir_id);
+  if (!import_ir.sem_ir) {
+    return;
+  }
 
-    SemIR::ImportIRId import_ir_id(import_index);
-    for (auto impl_index : llvm::seq(import_ir.sem_ir->impls().size())) {
-      SemIR::ImplId impl_id(impl_index);
+  for (auto impl_index : llvm::seq(import_ir.sem_ir->impls().size())) {
+    SemIR::ImplId impl_id(impl_index);
 
-      // Resolve the imported impl to a local impl ID.
-      ImportRefResolver resolver(context, import_ir_id);
-      resolver.Resolve(import_ir.sem_ir->impls().Get(impl_id).first_decl_id());
-    }
+    // Resolve the imported impl to a local impl ID.
+    ImportImpl(context, import_ir_id, impl_id);
   }
 }
 
+auto ImportImpl(Context& context, SemIR::ImportIRId import_ir_id,
+                SemIR::ImplId impl_id) -> void {
+  ImportRefResolver resolver(context, import_ir_id);
+  resolver.Resolve(context.import_irs()
+                       .Get(import_ir_id)
+                       .sem_ir->impls()
+                       .Get(impl_id)
+                       .first_decl_id());
+}
+
 }  // namespace Carbon::Check

toolchain/check/import_ref.h

@@ -42,8 +42,12 @@ auto VerifySameCanonicalImportIRInst(Context& context, SemIR::InstId prev_id,
 // ImportRefLoaded for use.
 auto LoadImportRef(Context& context, SemIR::InstId inst_id) -> void;
 
-// Load all impls declared in IRs imported into this context.
-auto ImportImpls(Context& context) -> void;
+// Load all impls declared in the api file corresponding to this impl file.
+auto ImportImplsFromApiFile(Context& context) -> void;
+
+// Load a specific impl declared in an imported IR.
+auto ImportImpl(Context& context, SemIR::ImportIRId import_ir_id,
+                SemIR::ImplId impl_id) -> void;
 
 }  // namespace Carbon::Check