Compare ImplWitnessAccess into Self as canonical constants (#5883)

This makes all `.Self` references in a facet type canonically the same
(which will remain true iff they refer to the same `Self` type in the
future), removing the need to do more complex comparisons between them
using the EntityName, interface, and index. This allows the comparison
of types containing `.Self` references to be done correctly regardless
of where the `.Self` appears, as such type expressions will all be
canonically equal if they otherwise equal now, regardless of whether
they are written in the context where `.Self` could have seen different
`Self` facet types.

In order to retain access to constraints on a base `.Self` facet type,
in the case of applying `where` to an existing facet type, we:
- Give the base facet type as a `RequirementBaseFacetType` constraint so
that eval of `WhereExpr` can find and copy all the constraints off of
it.
- Introduce eager/early rewrite constraint resolution, which allows a
constraint to eagerly resolve access to earlier rewrite constraints
(`where .A = () and .B = .A` is eagerly transformed into `where .A = ()
and .B = ()`) before the full constraint resolution step. This allows
use of rewrite constraints in larger type expressions, such as `where .A
= () and .B = C(.A)` and `C` will know that the argument is `()`.

Author: danakj

toolchain/check/context.h

@@ -229,6 +229,13 @@ class Context {
     return poisoned_concrete_impl_lookup_queries_;
   }
 
+  // A stack that tracks the rewrite constraints from a `where` expression being
+  // checked. The back of the stack is the currently checked `where` expression.
+  auto rewrites_stack()
+      -> llvm::SmallVector<Map<SemIR::ConstantId, SemIR::InstId>>& {
+    return rewrites_stack_;
+  }
+
   // --------------------------------------------------------------------------
   // Directly expose SemIR::File data accessors for brevity in calls.
   // --------------------------------------------------------------------------
@@ -423,6 +430,11 @@ class Context {
   // results at the end of the file. Any difference is diagnosed.
   llvm::SmallVector<PoisonedConcreteImplLookupQuery>
       poisoned_concrete_impl_lookup_queries_;
+
+  // A map from an ImplWitnessAccess on the LHS of a rewrite constraint to its
+  // value on the RHS. Used during checking of a `where` expression to allow
+  // constraints to access values from earlier constraints.
+  llvm::SmallVector<Map<SemIR::ConstantId, SemIR::InstId>> rewrites_stack_;
 };
 
 }  // namespace Carbon::Check

toolchain/check/eval.cpp

@@ -18,6 +18,7 @@
 #include "toolchain/check/facet_type.h"
 #include "toolchain/check/generic.h"
 #include "toolchain/check/import_ref.h"
+#include "toolchain/check/name_lookup.h"
 #include "toolchain/check/type.h"
 #include "toolchain/check/type_completion.h"
 #include "toolchain/diagnostics/diagnostic.h"
@@ -637,6 +638,13 @@ static auto GetConstantFacetTypeInfo(EvalContext& eval_context,
   // 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.
+  //
+  // This also implies that we may find `ImplWitnessAccessSubstituted`
+  // instructions in the LHS and RHS of these constraints, which are preserved
+  // to maintain them as an unresolved reference to an associated constant, but
+  // which must be handled gracefully during resolution. They will be replaced
+  // with the constant value of the `ImplWitnessAccess` below when they are
+  // substituted with a constant value.
   info.rewrite_constraints = orig.rewrite_constraints;
   if (!ResolveFacetTypeRewriteConstraints(eval_context.context(), loc_id,
                                           info.rewrite_constraints)) {
@@ -2049,31 +2057,32 @@ auto TryEvalTypedInst<SemIR::WhereExpr>(EvalContext& eval_context,
   auto typed_inst = inst.As<SemIR::WhereExpr>();
 
   Phase phase = Phase::Concrete;
-  SemIR::TypeId base_facet_type_id =
-      eval_context.GetTypeOfInst(typed_inst.period_self_id);
-  SemIR::Inst base_facet_inst =
-      eval_context.types().GetAsInst(base_facet_type_id);
   SemIR::FacetTypeInfo info = {.other_requirements = false};
 
-  // `where` provides that the base facet is an error, `type`, or a facet
-  // type.
-  if (auto facet_type = base_facet_inst.TryAs<SemIR::FacetType>()) {
-    info = eval_context.facet_types().Get(facet_type->facet_type_id);
-  } else if (base_facet_type_id == SemIR::ErrorInst::TypeId) {
-    return SemIR::ErrorInst::ConstantId;
-  } else {
-    CARBON_CHECK(base_facet_type_id == SemIR::TypeType::TypeId,
-                 "Unexpected type_id: {0}, inst: {1}", base_facet_type_id,
-                 base_facet_inst);
-  }
-
   // Add the constraints from the `WhereExpr` instruction into `info`.
   if (typed_inst.requirements_id.has_value()) {
     auto insts = eval_context.inst_blocks().Get(typed_inst.requirements_id);
     for (auto inst_id : insts) {
-      if (auto rewrite =
-              eval_context.insts().TryGetAs<SemIR::RequirementRewrite>(
+      if (auto base =
+              eval_context.insts().TryGetAs<SemIR::RequirementBaseFacetType>(
                   inst_id)) {
+        if (base->base_type_inst_id == SemIR::ErrorInst::TypeInstId) {
+          return SemIR::ErrorInst::ConstantId;
+        }
+
+        if (auto base_facet_type =
+                eval_context.insts().TryGetAs<SemIR::FacetType>(
+                    base->base_type_inst_id)) {
+          const auto& base_info =
+              eval_context.facet_types().Get(base_facet_type->facet_type_id);
+          info.extend_constraints.append(base_info.extend_constraints);
+          info.self_impls_constraints.append(base_info.self_impls_constraints);
+          info.rewrite_constraints.append(base_info.rewrite_constraints);
+          info.other_requirements |= base_info.other_requirements;
+        }
+      } else if (auto rewrite =
+                     eval_context.insts().TryGetAs<SemIR::RequirementRewrite>(
+                         inst_id)) {
         info.rewrite_constraints.push_back(
             {.lhs_id = rewrite->lhs_id, .rhs_id = rewrite->rhs_id});
       } else if (auto impls =

toolchain/check/eval_inst.cpp

@@ -333,6 +333,13 @@ auto EvalConstantInst(Context& context, SemIR::InstId inst_id,
   return ConstantEvalResult::NewSamePhase(inst);
 }
 
+auto EvalConstantInst(Context& context,
+                      SemIR::ImplWitnessAccessSubstituted inst)
+    -> ConstantEvalResult {
+  return ConstantEvalResult::Existing(
+      context.constant_values().Get(inst.value_id));
+}
+
 auto EvalConstantInst(Context& context,
                       SemIR::ImplWitnessAssociatedConstant inst)
     -> ConstantEvalResult {

toolchain/check/facet_type.cpp

@@ -24,9 +24,11 @@ namespace Carbon::Check {
 
 auto FacetTypeFromInterface(Context& context, SemIR::InterfaceId interface_id,
                             SemIR::SpecificId specific_id) -> SemIR::FacetType {
-  SemIR::FacetTypeId facet_type_id = context.facet_types().Add(
+  auto info =
       SemIR::FacetTypeInfo{.extend_constraints = {{interface_id, specific_id}},
-                           .other_requirements = false});
+                           .other_requirements = false};
+  info.Canonicalize();
+  SemIR::FacetTypeId facet_type_id = context.facet_types().Add(info);
   return {.type_id = SemIR::TypeType::TypeId, .facet_type_id = facet_type_id};
 }
 
@@ -58,6 +60,16 @@ static auto IncompleteFacetTypeDiagnosticBuilder(
   }
 }
 
+auto GetImplWitnessAccessWithoutSubstitution(Context& context,
+                                             SemIR::InstId inst_id)
+    -> SemIR::InstId {
+  if (auto inst = context.insts().TryGetAs<SemIR::ImplWitnessAccessSubstituted>(
+          inst_id)) {
+    return inst->impl_witness_access_id;
+  }
+  return inst_id;
+}
+
 auto InitialFacetTypeImplWitness(
     Context& context, SemIR::LocId witness_loc_id,
     SemIR::TypeInstId facet_type_inst_id, SemIR::TypeInstId self_type_inst_id,
@@ -123,8 +135,8 @@ auto InitialFacetTypeImplWitness(
   }
 
   for (auto rewrite : facet_type_info.rewrite_constraints) {
-    auto access =
-        context.insts().GetAs<SemIR::ImplWitnessAccess>(rewrite.lhs_id);
+    auto access = context.insts().GetAs<SemIR::ImplWitnessAccess>(
+        GetImplWitnessAccessWithoutSubstitution(context, rewrite.lhs_id));
     if (!WitnessQueryMatchesInterface(context, access.witness_id,
                                       interface_to_witness)) {
       continue;
@@ -252,6 +264,9 @@ auto AllocateFacetTypeImplWitness(Context& context,
 }
 
 namespace {
+// TODO: This class should go away, and we should just use the constant value of
+// the ImplWitnessAccess as a key in AccessRewriteValues, but that requires
+// changing its API to work with InstId instead of ImplWitnessAccess.
 struct FacetTypeConstraintValue {
   SemIR::EntityNameId entity_name_id;
   SemIR::ElementIndex access_index;
@@ -280,32 +295,6 @@ static auto GetFacetTypeConstraintValue(Context& context,
   return std::nullopt;
 }
 
-// Returns true if two values in a rewrite constraint are equivalent. Two
-// `ImplWitnessAccess` instructions that refer to the same associated constant
-// through the same facet value are treated as equivalent.
-static auto CompareFacetTypeConstraintValues(Context& context,
-                                             SemIR::InstId lhs_id,
-                                             SemIR::InstId rhs_id) -> bool {
-  if (lhs_id == rhs_id) {
-    return true;
-  }
-
-  auto lhs_access = context.insts().TryGetAs<SemIR::ImplWitnessAccess>(lhs_id);
-  auto rhs_access = context.insts().TryGetAs<SemIR::ImplWitnessAccess>(rhs_id);
-  if (lhs_access && rhs_access) {
-    auto lhs_access_value = GetFacetTypeConstraintValue(context, *lhs_access);
-    auto rhs_access_value = GetFacetTypeConstraintValue(context, *rhs_access);
-    // We do *not* want to get the evaluated result of `ImplWitnessAccess` here,
-    // we want to keep them as a reference to an associated constant for the
-    // resolution phase.
-    return lhs_access_value && rhs_access_value &&
-           *lhs_access_value == *rhs_access_value;
-  }
-
-  return context.constant_values().GetConstantInstId(lhs_id) ==
-         context.constant_values().GetConstantInstId(rhs_id);
-}
-
 // A mapping of each associated constant (represented as `ImplWitnessAccess`) to
 // its value (represented as an `InstId`). Used to track rewrite constraints,
 // with the LHS mapping to the resolved value of the RHS.
@@ -349,9 +338,10 @@ class AccessRewriteValues {
 
   auto SetFullyRewritten(Context& context, Value& value, SemIR::InstId inst_id)
       -> void {
-    CARBON_CHECK(
-        value.state == BeingRewritten ||
-        CompareFacetTypeConstraintValues(context, value.inst_id, inst_id));
+    if (value.state == FullyRewritten) {
+      CARBON_CHECK(context.constant_values().Get(value.inst_id) ==
+                   context.constant_values().Get(inst_id));
+    }
     value = {FullyRewritten, inst_id};
   }
 
@@ -459,13 +449,22 @@ class SubstImplWitnessAccessCallbacks : public SubstInstCallbacks {
         // There's no ImplWitnessAccess that we care about inside this
         // instruction.
         return SubstResult::FullySubstituted;
-      } else {
-        // SubstOperands will result in a Rebuild or ReuseUnchanged callback, so
-        // push the non-ImplWitnessAccess to get proper bracketing, allowing us
-        // to pop it in the paired callback.
+      }
+      if (auto subst =
+              context().insts().TryGetAs<SemIR::ImplWitnessAccessSubstituted>(
+                  rhs_inst_id)) {
+        // The reference to an associated constant was eagerly replaced with the
+        // value of an earlier rewrite constraint, but may need further
+        // substitution if it contains an `ImplWitnessAccess`.
+        rhs_inst_id = subst->value_id;
         substs_in_progress_.push_back(rhs_inst_id);
-        return SubstResult::SubstOperands;
+        return SubstResult::SubstAgain;
       }
+      // SubstOperands will result in a Rebuild or ReuseUnchanged callback, so
+      // push the non-ImplWitnessAccess to get proper bracketing, allowing us
+      // to pop it in the paired callback.
+      substs_in_progress_.push_back(rhs_inst_id);
+      return SubstResult::SubstOperands;
     }
 
     // If the access is going through a nested `ImplWitnessAccess`, that
@@ -581,8 +580,8 @@ auto ResolveFacetTypeRewriteConstraints(
   AccessRewriteValues rewrite_values;
 
   for (auto& constraint : rewrites) {
-    auto lhs_access =
-        context.insts().TryGetAs<SemIR::ImplWitnessAccess>(constraint.lhs_id);
+    auto lhs_access = context.insts().TryGetAs<SemIR::ImplWitnessAccess>(
+        GetImplWitnessAccessWithoutSubstitution(context, constraint.lhs_id));
     if (!lhs_access) {
       continue;
     }
@@ -591,8 +590,8 @@ auto ResolveFacetTypeRewriteConstraints(
   }
 
   for (auto& constraint : rewrites) {
-    auto lhs_access =
-        context.insts().TryGetAs<SemIR::ImplWitnessAccess>(constraint.lhs_id);
+    auto lhs_access = context.insts().TryGetAs<SemIR::ImplWitnessAccess>(
+        GetImplWitnessAccessWithoutSubstitution(context, constraint.lhs_id));
     if (!lhs_access) {
       continue;
     }
@@ -610,33 +609,40 @@ auto ResolveFacetTypeRewriteConstraints(
       return false;
     }
 
-    if (lhs_rewrite_value->state == AccessRewriteValues::FullyRewritten &&
-        !CompareFacetTypeConstraintValues(context, lhs_rewrite_value->inst_id,
-                                          rhs_subst_inst_id)) {
-      if (lhs_rewrite_value->inst_id != SemIR::ErrorInst::InstId) {
-        CARBON_DIAGNOSTIC(AssociatedConstantWithDifferentValues, Error,
-                          "associated constant {0} given two different "
-                          "values {1} and {2}",
-                          InstIdAsConstant, InstIdAsConstant, InstIdAsConstant);
-        // Use inst id ordering as a simple proxy for source ordering, to
-        // try to name the values in the same order they appear in the facet
-        // type.
-        auto source_order1 =
-            lhs_rewrite_value->inst_id.index < rhs_subst_inst_id.index
-                ? lhs_rewrite_value->inst_id
-                : rhs_subst_inst_id;
-        auto source_order2 =
-            lhs_rewrite_value->inst_id.index >= rhs_subst_inst_id.index
-                ? lhs_rewrite_value->inst_id
-                : rhs_subst_inst_id;
-        // TODO: It would be nice to note the places where the values are
-        // assigned but rewrite constraint instructions are from canonical
-        // constant values, and have no locations. We'd need to store a
-        // location along with them in the rewrite constraints.
-        context.emitter().Emit(loc_id, AssociatedConstantWithDifferentValues,
-                               constraint.lhs_id, source_order1, source_order2);
+    if (lhs_rewrite_value->state == AccessRewriteValues::FullyRewritten) {
+      auto rhs_existing_const_id =
+          context.constant_values().Get(lhs_rewrite_value->inst_id);
+      auto rhs_subst_const_id =
+          context.constant_values().Get(rhs_subst_inst_id);
+      if (rhs_subst_const_id != rhs_existing_const_id) {
+        if (rhs_existing_const_id != SemIR::ErrorInst::ConstantId) {
+          CARBON_DIAGNOSTIC(AssociatedConstantWithDifferentValues, Error,
+                            "associated constant {0} given two different "
+                            "values {1} and {2}",
+                            InstIdAsConstant, InstIdAsConstant,
+                            InstIdAsConstant);
+          // Use inst id ordering as a simple proxy for source ordering, to
+          // try to name the values in the same order they appear in the facet
+          // type.
+          auto source_order1 =
+              lhs_rewrite_value->inst_id.index < rhs_subst_inst_id.index
+                  ? lhs_rewrite_value->inst_id
+                  : rhs_subst_inst_id;
+          auto source_order2 =
+              lhs_rewrite_value->inst_id.index >= rhs_subst_inst_id.index
+                  ? lhs_rewrite_value->inst_id
+                  : rhs_subst_inst_id;
+          // TODO: It would be nice to note the places where the values are
+          // assigned but rewrite constraint instructions are from canonical
+          // constant values, and have no locations. We'd need to store a
+          // location along with them in the rewrite constraints.
+          context.emitter().Emit(loc_id, AssociatedConstantWithDifferentValues,
+                                 GetImplWitnessAccessWithoutSubstitution(
+                                     context, constraint.lhs_id),
+                                 source_order1, source_order2);
+        }
+        return false;
       }
-      return false;
     }
 
     rewrite_values.SetFullyRewritten(context, *lhs_rewrite_value,
@@ -652,8 +658,8 @@ auto ResolveFacetTypeRewriteConstraints(
   for (size_t i = 0; i < keep_size;) {
     auto& constraint = rewrites[i];
 
-    auto lhs_access =
-        context.insts().TryGetAs<SemIR::ImplWitnessAccess>(constraint.lhs_id);
+    auto lhs_access = context.insts().TryGetAs<SemIR::ImplWitnessAccess>(
+        GetImplWitnessAccessWithoutSubstitution(context, constraint.lhs_id));
     if (!lhs_access) {
       ++i;
       continue;

toolchain/check/facet_type.h

@@ -18,6 +18,16 @@ namespace Carbon::Check {
 auto FacetTypeFromInterface(Context& context, SemIR::InterfaceId interface_id,
                             SemIR::SpecificId specific_id) -> SemIR::FacetType;
 
+// Given an ImplWitnessAccessSubstituted, returns the InstId of the
+// ImplWitnessAccess. Otherwise, returns the input `inst_id` unchanged.
+//
+// This must be used when accessing the LHS of a rewrite constraint which has
+// not yet been resolved in order to preserve which associated constant is being
+// rewritten.
+auto GetImplWitnessAccessWithoutSubstitution(Context& context,
+                                             SemIR::InstId inst_id)
+    -> SemIR::InstId;
+
 // Creates a impl witness instruction for a facet type. The facet type is
 // required to be complete if `is_definition` is true or the facet type has
 // rewrites. Otherwise a placeholder witness is created, and