Resolve nested accesses in rewrite constraints (#5872)

A rewrite constraint like `.X = .Y.Z and .Y = .Self and .Z = ()` has a
nested `ImplWitnessAccess` `.Y.Z` (technically `(.Self.Y).Z`). The inner
access `.Self.Y` needs to be resolved (in this case to `.Self`) before
the outer `???.Z` can be resolved as `.Self.Z` which is `()`.

Author: danakj

toolchain/check/eval.cpp

@@ -635,10 +635,7 @@ static auto GetConstantFacetTypeInfo(EvalContext& eval_context,
   // Rewrite constraints are resolved first before replacing them with their
   // canonical instruction, so that in a `WhereExpr` we can work with the
   // `ImplWitnessAccess` references to `.Self` on the LHS of the constraints
-  // rather than the value of the associated constant they reference. It also
-  // ensures that any errors inserted during resolution will be seen by
-  // GetConstantValueIgnoringPeriodSelf() which will update the phase
-  // accordingly.
+  // rather than the value of the associated constant they reference.
   info.rewrite_constraints = orig.rewrite_constraints;
   if (!ResolveFacetTypeRewriteConstraints(eval_context.context(), loc_id,
                                           info.rewrite_constraints)) {

toolchain/check/facet_type.cpp

@@ -329,14 +329,18 @@ class AccessRewriteValues {
 
   auto InsertNotRewritten(Context& context, SemIR::ImplWitnessAccess access,
                           SemIR::InstId inst_id) -> void {
-    map_.insert({GetKey(context, access), {NotRewritten, inst_id}});
+    map_.insert({*GetKey(context, access), {NotRewritten, inst_id}});
   }
 
   // Finds and returns a pointer into the cache for a given ImplWitnessAccess.
   // The pointer will be invalidated by mutating the cache. Returns `nullptr`
   // if `access` is not found.
   auto FindRef(Context& context, SemIR::ImplWitnessAccess access) -> Value* {
-    auto it = map_.find(GetKey(context, access));
+    auto key = GetKey(context, access);
+    if (!key) {
+      return nullptr;
+    }
+    auto it = map_.find(*key);
     if (it == map_.end()) {
       return nullptr;
     }
@@ -391,8 +395,12 @@ class AccessRewriteValues {
     }
   };
 
-  auto GetKey(Context& context, SemIR::ImplWitnessAccess access) -> Key {
-    return {*GetFacetTypeConstraintValue(context, access)};
+  // Returns a key for the `access` to an associated context if the access is
+  // through a facet value. If the access it through another `ImplWitnessAccess`
+  // then no key is able to be made.
+  auto GetKey(Context& context, SemIR::ImplWitnessAccess access)
+      -> std::optional<Key> {
+    return GetFacetTypeConstraintValue(context, access);
   }
 
   // Try avoid heap allocations in the common case where there are a small
@@ -466,6 +474,18 @@ class SubstImplWitnessAccessCallbacks : public SubstInstCallbacks {
       }
     }
 
+    // If the access is going through a nested `ImplWitnessAccess`, that
+    // access needs to be resolved to a facet value first. If it can't be
+    // resolved then the outer one can not be either.
+    if (auto lookup = context().insts().TryGetAs<SemIR::LookupImplWitness>(
+            rhs_access->witness_id)) {
+      if (context().insts().Is<SemIR::ImplWitnessAccess>(
+              lookup->query_self_inst_id)) {
+        substs_in_progress_.push_back(rhs_inst_id);
+        return SubstResult::SubstOperandsAndRetry;
+      }
+    }
+
     auto* rewrite_value = rewrite_values_->FindRef(context(), *rhs_access);
     if (!rewrite_value) {
       // The RHS refers to an associated constant for which there is no rewrite
@@ -510,9 +530,9 @@ class SubstImplWitnessAccessCallbacks : public SubstInstCallbacks {
     auto subst_inst_id = substs_in_progress_.pop_back_val();
     if (auto access = context().insts().TryGetAs<SemIR::ImplWitnessAccess>(
             subst_inst_id)) {
-      auto* rewrite_value = rewrite_values_->FindRef(context(), *access);
-      CARBON_CHECK(rewrite_value);
-      rewrite_values_->SetFullyRewritten(context(), *rewrite_value, inst_id);
+      if (auto* rewrite_value = rewrite_values_->FindRef(context(), *access)) {
+        rewrite_values_->SetFullyRewritten(context(), *rewrite_value, inst_id);
+      }
     }
     return inst_id;
   }
@@ -521,10 +541,10 @@ class SubstImplWitnessAccessCallbacks : public SubstInstCallbacks {
     auto subst_inst_id = substs_in_progress_.pop_back_val();
     if (auto access = context().insts().TryGetAs<SemIR::ImplWitnessAccess>(
             subst_inst_id)) {
-      auto* rewrite_value = rewrite_values_->FindRef(context(), *access);
-      CARBON_CHECK(rewrite_value);
-      rewrite_values_->SetFullyRewritten(context(), *rewrite_value,
-                                         orig_inst_id);
+      if (auto* rewrite_value = rewrite_values_->FindRef(context(), *access)) {
+        rewrite_values_->SetFullyRewritten(context(), *rewrite_value,
+                                           orig_inst_id);
+      }
     }
     return orig_inst_id;
   }

toolchain/check/subst.cpp

@@ -312,26 +312,33 @@ auto SubstInst(Context& context, SemIR::InstId inst_id,
   while (index != -1) {
     auto& item = worklist[index];
 
-    if (item.is_repeated) {
-      // Pop the copy of the repeated item when we get back to the repeated
-      // item, and steal any work that was done there so we don't have to
-      // Subst() the repeated item again. The pop does not reallocate the
-      // worklist so does not invalidate `item`.
-      item.inst_id = worklist.Pop();
-    }
-
     if (item.is_expanded) {
       // Rebuild this item if necessary. Note that this might pop items from the
       // worklist but does not reallocate, so does not invalidate `item`.
-      item.inst_id = Rebuild(context, worklist, item.inst_id, callbacks);
-      index = item.next_index;
-      continue;
+      auto old_inst_id = std::exchange(
+          item.inst_id, Rebuild(context, worklist, item.inst_id, callbacks));
+      if (item.is_repeated && old_inst_id != item.inst_id) {
+        // SubstOperandsAndRetry was returned for the item, and the instruction
+        // was rebuilt from new operands, so go through Subst() again. Note that
+        // we've already called Rebuild so we don't want to leave this item as
+        // repeated, and call back to ReuseUnchanged for it again later unless
+        // the next call to Subst() asks for that.
+        item.is_expanded = false;
+        item.is_repeated = false;
+      } else {
+        index = item.next_index;
+        continue;
+      }
     }
 
     if (item.is_repeated) {
+      // SubstAgain was returned for the item, and the result of that Subst() is
+      // at the back of the worklist, which we pop. Note that popping from the
+      // worklist does not reallocate, so does not invalidate `item`.
+      //
       // When Subst returns SubstAgain, we must call back to Rebuild or
       // ReuseUnchanged for that work item.
-      item.inst_id = callbacks.ReuseUnchanged(item.inst_id);
+      item.inst_id = callbacks.ReuseUnchanged(worklist.Pop());
       index = item.next_index;
       continue;
     }
@@ -357,6 +364,9 @@ auto SubstInst(Context& context, SemIR::InstId inst_id,
       }
       case SubstInstCallbacks::SubstResult::SubstOperands:
         break;
+      case SubstInstCallbacks::SubstResult::SubstOperandsAndRetry:
+        item.is_repeated = true;
+        break;
     }
 
     // Extract the operands of this item into the worklist. Note that this

toolchain/check/subst.h

@@ -32,16 +32,28 @@ class SubstInstCallbacks {
     // Attempt to substitute again on the resulting instruction, acting like
     // recursion on the instruction itself.
     SubstAgain,
+    // Attempt to substitute into the operands of the instruction. If the InstId
+    // returned from Rebuild or ReuseUnchanged differs from the input (typically
+    // because some operand in the instruction changed), then the new
+    // instruction will be given to `Subst` again afterward. This allows for the
+    // uncommon case of substituting from the inside out.
+    SubstOperandsAndRetry,
   };
 
   // Performs any needed substitution into an instruction. The instruction ID
-  // should be updated as necessary to represent the new instruction. Returns
-  // FullySubstituted if the resulting instruction ID is fully-substituted.
-  // Return SubstOperands if substitution may be needed into operands of the
-  // instruction, or SubstAgain if the replaced instruction itself should have
-  // substitution applied to it again. When SubstOperands or SubstAgain is
-  // returned, it results in a call back to Rebuild or ReuseUnchanged when that
-  // instruction is done being substituted.
+  // should be updated as necessary to represent the new instruction.
+  //
+  // Return FullySubstituted if the resulting instruction ID is
+  // fully-substituted. Return SubstOperands if substitution may be needed into
+  // operands of the instruction, or SubstAgain if the replaced instruction
+  // itself should have substitution applied to it again. Return
+  // SubstOperandsAndRetry to recurse on the instructions operands and then
+  // substitute the resulting instruction afterward, if the instruction is
+  // replaced by a new one (typically due to Rebuild when the operands changed).
+  //
+  // When SubstOperands, SubstAgain, or SubstOperandsAndRetry is returned, it
+  // results in a call back to Rebuild or ReuseUnchanged when that instruction's
+  // substitution step is complete.
   virtual auto Subst(SemIR::InstId& inst_id) -> SubstResult = 0;
 
   // Rebuilds the type of an instruction from the substituted type instruction.

toolchain/check/testdata/facet/facet_assoc_const.carbon

@@ -601,6 +601,59 @@ fn J(T:! Z where .T0 = .T1 and .T1 = (.T2, .T3) and .T2 = .T4 and .T4 = .T5 and
   return ((), ());
 }
 
+// --- indirection_through_self_rhs.carbon
+library "[[@TEST_NAME]]";
+
+interface I {
+  let I1:! type;
+  let I2:! type;
+}
+
+interface J {
+  let J1:! I;
+}
+
+// The value of .I1 is (), but to know that requires resolving .J1 first then
+// .J1.I2.
+fn F(T:! I & J where .J1 = .Self and .I1 = .J1.I2 and .I2 = ()) -> T.I1 {
+  return ();
+}
+
+// --- indirection_through_not_self_rhs.carbon
+library "[[@TEST_NAME]]";
+
+interface I {
+  let I1:! type;
+  let I2:! type;
+}
+
+interface J {
+  let J1:! I;
+}
+
+// The value of .I1 is (), but to know that requires resolving .J1 first then
+// .J1.I2.
+fn F(U:! I where .I2 = (), T:! I & J where .J1 = U and .I1 = .J1.I2) -> T.I1 {
+  return ();
+}
+
+// --- indirection_through_unresolved_access_rhs.carbon
+library "[[@TEST_NAME]]";
+
+interface I {
+  let I1:! type;
+  let I2:! type;
+}
+
+interface J {
+  let J1:! I;
+}
+
+// If we assume the nested `.J1` access will resolve to a facet value, we may
+// loop forever trying to resolve the `.I2` access. We should gracefully accept
+// that it does not resolve further.
+fn F(T:! I & J where .I1 = .J1.I2) {}
+
 // CHECK:STDOUT: --- fail_cycle.carbon
 // CHECK:STDOUT:
 // CHECK:STDOUT: constants {

toolchain/check/testdata/facet/nested_facet_types.carbon

@@ -54,6 +54,21 @@ interface Z {
 
 fn F(FF:! ((Z where .T = .U and .U = .V) where .T = .U and .U = .V) where .T = .U and .U = .V) {}
 
+// --- fail_todo_nested_facet_types_same_associated_in_generic_parameter.carbon
+library "[[@TEST_NAME]]";
+
+interface Z {
+  let T:! type;
+  let U:! type;
+}
+class C(T:! type) {}
+
+// CHECK:STDERR: fail_todo_nested_facet_types_same_associated_in_generic_parameter.carbon:[[@LINE+4]]:12: error: associated constant `.(Z.T)` given two different values `C(.(Z.U))` and `C(.(Z.U))` [AssociatedConstantWithDifferentValues]
+// CHECK:STDERR: fn F(FF:! ((Z where .T = C(.U)) where .T = C(.U)) where .T = C(.U)) {}
+// CHECK:STDERR:            ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// CHECK:STDERR:
+fn F(FF:! ((Z where .T = C(.U)) where .T = C(.U)) where .T = C(.U)) {}
+
 // --- fail_nested_facet_types_different.carbon
 library "[[@TEST_NAME]]";
 
@@ -81,6 +96,21 @@ interface Z {
 // CHECK:STDERR:
 fn F(FF:! (Z where .T = .U) where .T = {}) {}
 
+// --- fail_nested_facet_types_different_with_associated_in_generic_parameter.carbon
+library "[[@TEST_NAME]]";
+
+interface Z {
+  let T:! type;
+  let U:! type;
+}
+class C(T:! type) {}
+
+// CHECK:STDERR: fail_nested_facet_types_different_with_associated_in_generic_parameter.carbon:[[@LINE+4]]:11: error: associated constant `.(Z.T)` given two different values `C(.(Z.U))` and `{}` [AssociatedConstantWithDifferentValues]
+// CHECK:STDERR: fn F(FF:! (Z where .T = C(.U)) where .T = {}) {}
+// CHECK:STDERR:           ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// CHECK:STDERR:
+fn F(FF:! (Z where .T = C(.U)) where .T = {}) {}
+
 // --- nested_facet_types_same_with_bitand.carbon
 library "[[@TEST_NAME]]";