Factor out GetCanonicalFileAndInstId for code sharing. (#5362)

Co-authored-by: Richard Smith <richard@metafoo.co.uk>

toolchain/check/import_ref.cpp

@@ -112,41 +112,21 @@ static auto AddLoadedImportRef(Context& context, SemIR::TypeId type_id,
 }
 
 static auto GetCanonicalImportIRInst(Context& context,
-                                     const SemIR::File* cursor_ir,
-                                     SemIR::InstId cursor_inst_id)
+                                     const SemIR::File* target_ir,
+                                     SemIR::InstId target_inst_id)
     -> SemIR::ImportIRInst {
-  while (true) {
-    // Step through an instruction with an imported location to the imported
-    // instruction.
-    auto loc_id = cursor_ir->insts().GetLocId(cursor_inst_id);
-    if (loc_id.kind() == SemIR::LocId::Kind::ImportIRInstId) {
-      auto import_ir_inst =
-          cursor_ir->import_ir_insts().Get(loc_id.import_ir_inst_id());
-      cursor_ir = cursor_ir->import_irs().Get(import_ir_inst.ir_id()).sem_ir;
-      cursor_inst_id = import_ir_inst.inst_id();
-      continue;
-    }
-
-    // Step through export declarations to their exported value.
-    if (auto export_decl =
-            cursor_ir->insts().TryGetAs<SemIR::ExportDecl>(cursor_inst_id)) {
-      cursor_inst_id = export_decl->value_id;
-      continue;
-    }
-
-    // Reached a non-imported entity.
-    break;
-  }
+  auto [canonical_ir, canonical_inst_id] =
+      GetCanonicalFileAndInstId(target_ir, target_inst_id);
 
   auto ir_id = SemIR::ImportIRId::None;
-  if (cursor_ir != &context.sem_ir()) {
+  if (canonical_ir != &context.sem_ir()) {
     // This uses AddImportIR in case it was indirectly found, which can
     // happen with two or more steps of exports.
     ir_id = AddImportIR(context, {.decl_id = SemIR::InstId::None,
                                   .is_export = false,
-                                  .sem_ir = cursor_ir});
+                                  .sem_ir = canonical_ir});
   }
-  return SemIR::ImportIRInst(ir_id, cursor_inst_id);
+  return SemIR::ImportIRInst(ir_id, canonical_inst_id);
 }
 
 auto GetCanonicalImportIRInst(Context& context, SemIR::InstId inst_id)

toolchain/check/import_ref.h

@@ -26,7 +26,9 @@ auto AddImportRef(Context& context, SemIR::ImportIRInst import_ir_inst,
                   SemIR::EntityNameId entity_name_id) -> SemIR::InstId;
 
 // Returns the canonical IR inst for an entity. Returns an `ImportIRInst` with
-// a `None` ir_id for an entity that was not imported.
+// a `None` ir_id for an entity that was not imported. This may add an
+// `ImportIR` if it finds an indirectly referenced `File` that hasn't previously
+// been found.
 auto GetCanonicalImportIRInst(Context& context, SemIR::InstId inst_id)
     -> SemIR::ImportIRInst;
 

toolchain/driver/compile_subcommand.cpp

@@ -25,6 +25,7 @@
 #include "toolchain/parse/parse.h"
 #include "toolchain/parse/tree_and_subtrees.h"
 #include "toolchain/sem_ir/formatter.h"
+#include "toolchain/sem_ir/import_ir.h"
 #include "toolchain/sem_ir/inst_namer.h"
 #include "toolchain/source/source_buffer.h"
 
@@ -562,25 +563,9 @@ auto CompilationUnit::PostCheck() -> void {
   if (vlog_stream_ || print) {
     // Omit entities imported from files that we are not dumping.
     auto should_format_entity = [&](SemIR::InstId entity_inst_id) -> bool {
-      // TODO: Reuse `GetCanonicalImportIRInst`. Currently it depends on
-      // `Check::Context`, which we don't have access to here.
-      const SemIR::File* file = &*sem_ir_;
-      while (true) {
-        auto loc_id = file->insts().GetLocId(entity_inst_id);
-        if (loc_id.kind() != SemIR::LocId::Kind::ImportIRInstId) {
-          return true;
-        }
-        auto import_ir_inst =
-            file->import_ir_insts().Get(loc_id.import_ir_inst_id());
-        const auto* import_file =
-            file->import_irs().Get(import_ir_inst.ir_id()).sem_ir;
-        CARBON_CHECK(import_file);
-        if (!IncludeInDumps(import_file->filename())) {
-          return false;
-        }
-        file = import_file;
-        entity_inst_id = import_ir_inst.inst_id();
-      }
+      auto [import_ir, _] =
+          SemIR::GetCanonicalFileAndInstId(&*sem_ir_, entity_inst_id);
+      return IncludeInDumps(import_ir->filename());
     };
 
     SemIR::Formatter formatter(&*sem_ir_, should_format_entity);

toolchain/sem_ir/BUILD

@@ -85,6 +85,7 @@ cc_library(
         "function.cpp",
         "generic.cpp",
         "impl.cpp",
+        "import_ir.cpp",
         "name.cpp",
         "name_scope.cpp",
         "pattern.cpp",

toolchain/sem_ir/import_ir.cpp

@@ -0,0 +1,54 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "toolchain/sem_ir/import_ir.h"
+
+#include "toolchain/sem_ir/file.h"
+
+namespace Carbon::SemIR {
+
+auto ImportIR::Print(llvm::raw_ostream& out) const -> void {
+  out << "{decl_id: " << decl_id
+      << ", is_export: " << (is_export ? "true" : "false") << "}";
+}
+
+auto ImportIRInst::Print(llvm::raw_ostream& out) const -> void {
+  out << "{ir_id: " << ir_id() << ", ";
+  if (ir_id() == ImportIRId::Cpp) {
+    out << "clang_source_loc_id: " << clang_source_loc_id();
+  } else {
+    out << "inst_id: " << inst_id();
+  }
+  out << "}";
+}
+
+auto GetCanonicalFileAndInstId(const File* sem_ir, SemIR::InstId inst_id)
+    -> std::pair<const File*, InstId> {
+  while (true) {
+    // Step through an instruction with an imported location to the imported
+    // instruction.
+    if (auto loc_id = sem_ir->insts().GetLocId(inst_id);
+        loc_id.kind() == SemIR::LocId::Kind::ImportIRInstId) {
+      auto import_ir_inst =
+          sem_ir->import_ir_insts().Get(loc_id.import_ir_inst_id());
+      sem_ir = sem_ir->import_irs().Get(import_ir_inst.ir_id()).sem_ir;
+      inst_id = import_ir_inst.inst_id();
+      continue;
+    }
+
+    // Step through export declarations to their exported value.
+    if (auto export_decl =
+            sem_ir->insts().TryGetAs<SemIR::ExportDecl>(inst_id)) {
+      inst_id = export_decl->value_id;
+      continue;
+    }
+
+    // Reached a non-imported entity.
+    break;
+  }
+
+  return std::make_pair(sem_ir, inst_id);
+}
+
+}  // namespace Carbon::SemIR

toolchain/sem_ir/import_ir.h

@@ -13,10 +13,7 @@ namespace Carbon::SemIR {
 
 // A reference to an imported IR.
 struct ImportIR : public Printable<ImportIR> {
-  auto Print(llvm::raw_ostream& out) const -> void {
-    out << "{decl_id: " << decl_id
-        << ", is_export: " << (is_export ? "true" : "false") << "}";
-  }
+  auto Print(llvm::raw_ostream& out) const -> void;
 
   // The `import` declaration.
   InstId decl_id;
@@ -42,15 +39,7 @@ class ImportIRInst : public Printable<ImportIRInst> {
   explicit ImportIRInst(ClangSourceLocId clang_source_loc_id)
       : ir_id_(ImportIRId::Cpp), clang_source_loc_id_(clang_source_loc_id) {}
 
-  auto Print(llvm::raw_ostream& out) const -> void {
-    out << "{ir_id: " << ir_id() << ", ";
-    if (ir_id() == ImportIRId::Cpp) {
-      out << "clang_source_loc_id: " << clang_source_loc_id();
-    } else {
-      out << "inst_id: " << inst_id();
-    }
-    out << "}";
-  }
+  auto Print(llvm::raw_ostream& out) const -> void;
 
   friend auto operator==(const ImportIRInst& lhs, const ImportIRInst& rhs)
       -> bool {
@@ -81,6 +70,12 @@ class ImportIRInst : public Printable<ImportIRInst> {
   };
 };
 
+// Returns the canonical `File` and `InstId` for an entity, tracing imported
+// instructions. Note the returned `File` might not be directly imported by the
+// input `sem_ir`.
+auto GetCanonicalFileAndInstId(const File* sem_ir, SemIR::InstId inst_id)
+    -> std::pair<const File*, InstId>;
+
 }  // namespace Carbon::SemIR
 
 #endif  // CARBON_TOOLCHAIN_SEM_IR_IMPORT_IR_H_