Check if context belongs to inference table before using.

PiperOrigin-RevId: 750581701

Author: erinzmoore

xls/dslx/type_system_v2/inference_table.cc

@@ -400,17 +400,16 @@ class InferenceTableImpl : public InferenceTable {
         (it == type_annotations_per_type_variable_.end())
             ? std::vector<const TypeAnnotation*>()
             : it->second;
-    if (parametric_context.has_value()) {
-      if (mutable_parametric_context_data_.contains(*parametric_context)) {
-        const auto& invocation_specific_annotations =
-            mutable_parametric_context_data_.at(*parametric_context)
-                .type_annotations_per_type_variable;
-        const auto invocation_specific_it =
-            invocation_specific_annotations.find(variable);
-        if (invocation_specific_it != invocation_specific_annotations.end()) {
-          absl::c_copy(invocation_specific_it->second,
-                       std::back_inserter(result));
-        }
+    if (parametric_context.has_value() &&
+        (*parametric_context)->node()->owner() == &module_) {
+      const auto& invocation_specific_annotations =
+          mutable_parametric_context_data_.at(*parametric_context)
+              .type_annotations_per_type_variable;
+      const auto invocation_specific_it =
+          invocation_specific_annotations.find(variable);
+      if (invocation_specific_it != invocation_specific_annotations.end()) {
+        absl::c_copy(invocation_specific_it->second,
+                     std::back_inserter(result));
       }
     }
     return result;
@@ -604,7 +603,7 @@ class InferenceTableImpl : public InferenceTable {
       std::optional<const ParametricContext*> context,
       const InferenceVariable* variable, const TypeAnnotation* annotation) {
     CHECK(variable->kind() == InferenceVariableKind::kType);
-    if (context.has_value()) {
+    if (context.has_value() && (*context)->node()->owner() == &module_) {
       mutable_parametric_context_data_.at(*context)
           .type_annotations_per_type_variable[variable]
           .push_back(annotation);

xls/dslx/type_system_v2/typecheck_module_v2_test.cc

@@ -6200,6 +6200,53 @@ fn main() -> u5 {
               IsOkAndHolds(HasTypeInfo(HasNodeWithType("var", "uN[5]"))));
 }
 
+TEST(TypecheckV2Test, ImportParametricFunctionWithMultipleInvocations) {
+  constexpr std::string_view kImported = R"(
+pub fn add_one(x: u32) -> u32 { x + 1 }
+
+pub fn some_function<N: u32, M: u32 = { add_one(N) }>() -> uN[M] { uN[M]:0 }
+
+pub fn another_fn() -> u3 { some_function<2>() }
+
+pub fn parametric_call<M: u32>() -> uN[M] { some_function<3, M>() }
+
+)";
+  constexpr std::string_view kInt = R"(
+import imported;
+
+pub fn mid() -> u32 { imported::some_function<31>() }
+
+pub fn default_import<N: u32, M: u32 = { imported::add_one(N) }>() -> uN[M] { uN[M]:0 }
+
+)";
+  constexpr std::string_view kProgram = R"(
+import imported;
+import int;
+
+const VAR = imported::some_function<4>();
+const VAR2 = imported::some_function<3, 6>();
+const VAR3 = imported::some_function<7>();
+const VAR4 = imported::another_fn();
+const VAR5 = imported::parametric_call<15>();
+const VAR6 = int::mid();
+
+fn main() -> u26 {
+   int::default_import<25>()
+}
+)";
+  ImportData import_data = CreateImportDataForTest();
+  XLS_EXPECT_OK(TypecheckV2(kImported, "imported", &import_data).status());
+  XLS_EXPECT_OK(TypecheckV2(kInt, "int", &import_data).status());
+  EXPECT_THAT(
+      TypecheckV2(kProgram, "main", &import_data),
+      IsOkAndHolds(HasTypeInfo(AllOf(
+          HasNodeWithType("VAR", "uN[5]"), HasNodeWithType("VAR2", "uN[6]"),
+          HasNodeWithType("VAR4", "uN[3]"), HasNodeWithType("VAR5", "uN[15]"),
+          HasNodeWithType("VAR6", "uN[32]"),
+          HasNodeWithType("int::default_import<25>()", "uN[26]"),
+          HasNodeWithType("VAR3", "uN[8]")))));
+}
+
 TEST(TypecheckV2Test, ImportParametricFunction) {
   constexpr std::string_view kImported = R"(
 pub fn some_function<N: u32>() -> uN[N] { uN[N]:0 }