Impls in final match_first block aren't always final (#7512)

If an earlier impl may match a more specific query, then later impls can
not be treated as final for the given query.

See
https://github.com/carbon-language/carbon-lang/blob/de8b03faa3178ae683d8e7124fbcba81eb88e00c/proposals/p005337-interface-extension-and-final-impl-update.md#using-associated-constants-from-impls-in-a-final-match_first

toolchain/check/impl_lookup.cpp

@@ -702,6 +702,12 @@ static auto CollectCandidateImplsForQuery(
     }
   }
 
+  // For each `match_first` block, track the position of the first impl that is
+  // more specific than the query, and could thus match the query once it is
+  // made more specific. Any impls with a later position in the same
+  // `match_first` block cannot be treated as final.
+  Map<SemIR::InstId, int, 16> first_more_specific_impl_in_match_first;
+
   for (auto [id, impl] : context.impls().enumerate()) {
     CARBON_CHECK(impl.witness_id.has_value());
 
@@ -711,16 +717,9 @@ static auto CollectCandidateImplsForQuery(
       continue;
     }
 
-    if (final_only) {
-      if (!TreatImplAsFinal(context, impl) &&
-          // TODO: For `match_first_is_final`, the impl can only be treated as
-          // final if there's no `impl` above it in the match_first block that
-          // overlaps with the query (such that a more specific query might
-          // choose it). See
-          // https://github.com/carbon-language/carbon-lang/blob/de8b03faa3178ae683d8e7124fbcba81eb88e00c/proposals/p005337-interface-extension-and-final-impl-update.md#using-associated-constants-from-impls-in-a-final-match_first
-          !impl.match_first_is_final) {
-        continue;
-      }
+    if (final_only && !TreatImplAsFinal(context, impl) &&
+        !impl.match_first_is_final) {
+      continue;
     }
 
     if (llvm::is_contained(context.forbidden_impls(), id)) {
@@ -747,12 +746,30 @@ static auto CollectCandidateImplsForQuery(
     if (!type_structure) {
       continue;
     }
+
     // TODO: We can skip the comparison here if the `impl_interface_const_id` is
     // not symbolic, since when the interface and specific ids match, and they
     // aren't symbolic, the structure will be identical.
     if (!query_type_structure.CompareStructure(
             TypeStructure::CompareTest::IsEqualToOrMoreSpecificThan,
             *type_structure)) {
+      // If the query does not match this impl, but the impl is part of a final
+      // `match_first` block, then we also check to see if the impl is _more
+      // specific_ than the query. Meaning that the query, once specialized,
+      // could match the impl. In that case, impls that come after can not be
+      // treated as final.
+      if (final_only && impl.match_first_is_final) {
+        if (type_structure->CompareStructure(
+                TypeStructure::CompareTest::IsEqualToOrMoreSpecificThan,
+                query_type_structure)) {
+          auto result = first_more_specific_impl_in_match_first.Insert(
+              impl.match_first_id, impl.match_first_position);
+          if (!result.is_inserted()) {
+            result.value() =
+                std::min(result.value(), impl.match_first_position);
+          }
+        }
+      }
       continue;
     }
 
@@ -761,6 +778,25 @@ static auto CollectCandidateImplsForQuery(
                                 impl.match_first_position});
   }
 
+  if (final_only) {
+    // When searching for `final_only`: Remove candidates if they are in a final
+    // `match_first` block, and there is a more specific impl in an earlier
+    // position in that block. That prevents the later impl from being
+    // considered final.
+    llvm::erase_if(candidates.impls, [&](auto& candidate) {
+      if (!candidate.match_first_block.has_value()) {
+        return false;
+      }
+      auto result = first_more_specific_impl_in_match_first.Lookup(
+          candidate.match_first_block);
+      if (!result) {
+        return false;
+      }
+      int first_more_specific_position = result.value();
+      return candidate.match_first_position > first_more_specific_position;
+    });
+  }
+
   auto compare = [](auto& lhs, auto& rhs) -> bool {
     // If they are in the same block, then order wins. Final impls will always
     // be in the same block if they overlap, as will impls that have the same

toolchain/check/testdata/match_first/prioritization.carbon

@@ -355,3 +355,120 @@ final match_first {
 fn F[U: X]() {
   {.y = type} as ().(Z(C(U)).Z1);
 }
+
+// --- fail_non_final_match_first_is_not_final.carbon
+library "[[@TEST_NAME]]";
+
+interface Z(T: type) { let Z1: type; }
+interface Y {}
+
+class C(T: type);
+
+impl forall [T: Y] () as Z(C(T)) where .Z1 = {.y: type} {}
+
+match_first {
+  impl forall [T: Y] () as Z(C(T)) where .Z1 = {.y: type};
+}
+
+fn F(generic T: Y) {
+  // This query is not concrete, so the impl will not give a final witness, so
+  // we can't use the rewrite constraint's RHS here.
+  // CHECK:STDERR: fail_non_final_match_first_is_not_final.carbon:[[@LINE+4]]:3: error: `Core.As` implicitly referenced here, but package `Core` not found [CoreNotFound]
+  // CHECK:STDERR:   {.y = type} as ().(Z(C(T)).Z1);
+  // CHECK:STDERR:   ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+  // CHECK:STDERR:
+  {.y = type} as ().(Z(C(T)).Z1);
+}
+
+// --- final_not_overlapping_query_in_match_first.carbon
+library "[[@TEST_NAME]]";
+
+interface Z(T: type) { let Z1: type; }
+interface W {}
+
+class C(T: type);
+class D(T: type);
+
+final match_first {
+  // Arbitrary impl at the start of the block that doesn't overlap anything.
+  impl C({}) as W {}
+
+  // A more specific impl comes first.
+  impl forall [T: type] () as Z(C(T)) where .Z1 = {.c: type} {}
+  // Another different but more specific impl comes next.
+  impl forall [T: type] () as Z(D(T)) where .Z1 = {.d: type} {}
+  // A more general impl comes last. While this catches more things, if T
+  // _could_ be specialized as `C(T)` or `D(T)` in the future, we can't use this
+  // impl as final.
+  impl forall [T: type] () as Z(T) where .Z1 = {.t: type} {}
+}
+
+fn F() {
+  // The first matching impl in the match_first is also the most specific
+  // overlapping impl, so it's considered final.
+  {.c = type} as ().(Z(C({})).Z1);
+
+  // The first matching impl in the match_first is also the most specific
+  // overlapping impl, so it's considered final.
+  {.d = type} as ().(Z(D({})).Z1);
+
+  // Since this can't match C(T) or D(T), the first matching impl in the
+  // match_first is also the most specific overlapping impl, so it's considered
+  // final.
+  class E;
+  {.t = type} as ().(Z(E).Z1);
+}
+
+// --- fail_final_overlapping_query_in_match_first.carbon
+library "[[@TEST_NAME]]";
+
+interface Z(T: type) { let Z1: type; }
+
+class C(T: type);
+
+final match_first {
+  // A more specific impl comes first.
+  impl forall [T: type] () as Z(C(T)) where .Z1 = {.c: type} {}
+  // A more general impl comes last. While this catches more things, if T
+  // _could_ be specialized as `C(T)` in the future, we can't use this impl as
+  // final.
+  impl forall [T: type] () as Z(T) where .Z1 = {.t: type} {}
+}
+
+fn F(generic T: type) {
+  // Could match C(T) later, but doesn't yet. So it's not considered final.
+  // CHECK:STDERR: fail_final_overlapping_query_in_match_first.carbon:[[@LINE+4]]:3: error: `Core.As` implicitly referenced here, but package `Core` not found [CoreNotFound]
+  // CHECK:STDERR:   {.t = type} as ().(Z(T).Z1);
+  // CHECK:STDERR:   ^~~~~~~~~~~~~~~~~~~~~~~~~~~
+  // CHECK:STDERR:
+  {.t = type} as ().(Z(T).Z1);
+}
+
+// --- fail_final_overlapping_query_in_match_first_reverse_order.carbon
+library "[[@TEST_NAME]]";
+
+interface Z(T: type) { let Z1: type; }
+
+class C(T: type);
+
+// The match_first is in the opposite order as the impls are introduced.
+impl forall [T: type] () as Z(T) where .Z1 = {.t: type} {}
+impl forall [T: type] () as Z(C(T)) where .Z1 = {.c: type} {}
+
+final match_first {
+  // A more specific impl comes first.
+  impl forall [T: type] () as Z(C(T)) where .Z1 = {.c: type};
+  // A more general impl comes last. While this catches more things, if T
+  // _could_ be specialized as `C(T)` in the future, we can't use this impl as
+  // final.
+  impl forall [T: type] () as Z(T) where .Z1 = {.t: type};
+}
+
+fn F(generic T: type) {
+  // Could match C(T) later, but doesn't yet. So it's not considered final.
+  // CHECK:STDERR: fail_final_overlapping_query_in_match_first_reverse_order.carbon:[[@LINE+4]]:3: error: `Core.As` implicitly referenced here, but package `Core` not found [CoreNotFound]
+  // CHECK:STDERR:   {.t = type} as ().(Z(T).Z1);
+  // CHECK:STDERR:   ^~~~~~~~~~~~~~~~~~~~~~~~~~~
+  // CHECK:STDERR:
+  {.t = type} as ().(Z(T).Z1);
+}