Support extended scopes that are parameterized types (#4524)

* The `extended_scopes` in a `NameScope` were represented by a
`NameScopeId`. Replace that with an `InstId` of an instruction returning
the type that is extending this name scope.
* `Context::LookupQualifiedName` now can take multiple scopes to look
in.
* `GetAsLookupScope` was moved out of `member_access.cpp` and is now
`Context::AppendLookupScopesForConstant`

This PR also fixes some existing issues that were revealed as part of
writing and testing this PR:
* Additional validation and handling of invalid ids.
* `extend impl` in a class is not properly imported yet, but at least
now it doesn't crash.

The change to use an `InstId` also allowed some diagnostics and
formatting to be improved.

---------

Co-authored-by: Josh L <[email protected]>

Author: josh11b

toolchain/check/context.cpp

@@ -321,10 +321,11 @@ auto Context::LookupUnqualifiedName(Parse::NodeId node_id,
   // Walk the non-lexical scopes and perform lookups into each of them.
   for (auto [index, lookup_scope_id, specific_id] :
        llvm::reverse(non_lexical_scopes)) {
-    if (auto non_lexical_result = LookupQualifiedName(
-            node_id, name_id,
-            {.name_scope_id = lookup_scope_id, .specific_id = specific_id},
-            /*required=*/false);
+    if (auto non_lexical_result =
+            LookupQualifiedName(node_id, name_id,
+                                LookupScope{.name_scope_id = lookup_scope_id,
+                                            .specific_id = specific_id},
+                                /*required=*/false);
         non_lexical_result.inst_id.is_valid()) {
       return non_lexical_result;
     }
@@ -440,11 +441,59 @@ struct ProhibitedAccessInfo {
   bool is_parent_access;
 };
 
+auto Context::AppendLookupScopesForConstant(
+    SemIRLoc loc, SemIR::ConstantId base_const_id,
+    llvm::SmallVector<LookupScope>* scopes) -> bool {
+  auto base_id = constant_values().GetInstId(base_const_id);
+  auto base = insts().Get(base_id);
+  if (auto base_as_namespace = base.TryAs<SemIR::Namespace>()) {
+    scopes->push_back(
+        LookupScope{.name_scope_id = base_as_namespace->name_scope_id,
+                    .specific_id = SemIR::SpecificId::Invalid});
+    return true;
+  }
+  if (auto base_as_class = base.TryAs<SemIR::ClassType>()) {
+    TryToDefineType(GetTypeIdForTypeConstant(base_const_id), [&] {
+      CARBON_DIAGNOSTIC(QualifiedExprInIncompleteClassScope, Error,
+                        "member access into incomplete class {0}",
+                        InstIdAsType);
+      return emitter().Build(loc, QualifiedExprInIncompleteClassScope, base_id);
+    });
+    auto& class_info = classes().Get(base_as_class->class_id);
+    scopes->push_back(LookupScope{.name_scope_id = class_info.scope_id,
+                                  .specific_id = base_as_class->specific_id});
+    return true;
+  }
+  if (auto base_as_facet_type = base.TryAs<SemIR::FacetType>()) {
+    TryToDefineType(GetTypeIdForTypeConstant(base_const_id), [&] {
+      CARBON_DIAGNOSTIC(QualifiedExprInUndefinedInterfaceScope, Error,
+                        "member access into undefined interface {0}",
+                        InstIdAsType);
+      return emitter().Build(loc, QualifiedExprInUndefinedInterfaceScope,
+                             base_id);
+    });
+    const auto& facet_type_info =
+        facet_types().Get(base_as_facet_type->facet_type_id);
+    for (auto interface : facet_type_info.impls_constraints) {
+      auto& interface_info = interfaces().Get(interface.interface_id);
+      scopes->push_back(LookupScope{.name_scope_id = interface_info.scope_id,
+                                    .specific_id = interface.specific_id});
+    }
+    return true;
+  }
+  // TODO: Per the design, if `base_id` is any kind of type, then lookup should
+  // treat it as a name scope, even if it doesn't have members. For example,
+  // `(i32*).X` should fail because there's no name `X` in `i32*`, not because
+  // there's no name `X` in `type`.
+  return false;
+}
+
 auto Context::LookupQualifiedName(SemIRLoc loc, SemIR::NameId name_id,
-                                  LookupScope scope, bool required,
+                                  llvm::ArrayRef<LookupScope> lookup_scopes,
+                                  bool required,
                                   std::optional<AccessInfo> access_info)
     -> LookupResult {
-  llvm::SmallVector<LookupScope> scopes = {scope};
+  llvm::SmallVector<LookupScope> scopes(lookup_scopes);
 
   // TODO: Support reporting of multiple prohibited access.
   llvm::SmallVector<ProhibitedAccessInfo> prohibited_accesses;
@@ -457,6 +506,10 @@ auto Context::LookupQualifiedName(SemIRLoc loc, SemIR::NameId name_id,
   // Walk this scope and, if nothing is found here, the scopes it extends.
   while (!scopes.empty()) {
     auto [scope_id, specific_id] = scopes.pop_back_val();
+    if (!scope_id.is_valid()) {
+      has_error = true;
+      continue;
+    }
     const auto& name_scope = name_scopes().Get(scope_id);
     has_error |= name_scope.has_error;
 
@@ -479,13 +532,25 @@ auto Context::LookupQualifiedName(SemIRLoc loc, SemIR::NameId name_id,
     if (!scope_result_id.is_valid() || is_access_prohibited) {
       // If nothing is found in this scope or if we encountered an invalid
       // access, look in its extended scopes.
-      auto extended = name_scope.extended_scopes;
+      const auto& extended = name_scope.extended_scopes;
       scopes.reserve(scopes.size() + extended.size());
       for (auto extended_id : llvm::reverse(extended)) {
-        // TODO: Track a constant describing the extended scope, and substitute
-        // into it to determine its corresponding specific.
-        scopes.push_back({.name_scope_id = extended_id,
-                          .specific_id = SemIR::SpecificId::Invalid});
+        // Substitute into the constant describing the extended scope to
+        // determine its corresponding specific.
+        CARBON_CHECK(extended_id.is_valid());
+        SemIR::ConstantId const_id =
+            GetConstantValueInSpecific(sem_ir(), specific_id, extended_id);
+
+        DiagnosticAnnotationScope annotate_diagnostics(
+            &emitter(), [&](auto& builder) {
+              CARBON_DIAGNOSTIC(FromExtendHere, Note,
+                                "declared as an extended scope here");
+              builder.Note(extended_id, FromExtendHere);
+            });
+        if (!AppendLookupScopesForConstant(loc, const_id, &scopes)) {
+          // TODO: Handle case where we have a symbolic type and instead should
+          // look in its type.
+        }
       }
       is_parent_access |= !extended.empty();
       continue;

toolchain/check/context.h

@@ -225,10 +225,19 @@ class Context {
                               const SemIR::NameScope& scope)
       -> std::pair<SemIR::InstId, SemIR::AccessKind>;
 
-  // Performs a qualified name lookup in a specified scope and in scopes that
-  // it extends, returning the referenced instruction.
+  // Appends the lookup scopes corresponding to `base_const_id` to `*scopes`.
+  // Returns `false` if not a scope. On invalid scopes, prints a diagnostic, but
+  // still updates `*scopes` and returns `true`.
+  auto AppendLookupScopesForConstant(SemIRLoc loc,
+                                     SemIR::ConstantId base_const_id,
+                                     llvm::SmallVector<LookupScope>* scopes)
+      -> bool;
+
+  // Performs a qualified name lookup in a specified scopes and in scopes that
+  // they extend, returning the referenced instruction.
   auto LookupQualifiedName(SemIRLoc loc, SemIR::NameId name_id,
-                           LookupScope scope, bool required = true,
+                           llvm::ArrayRef<LookupScope> lookup_scopes,
+                           bool required = true,
                            std::optional<AccessInfo> access_info = std::nullopt)
       -> LookupResult;
 

toolchain/check/deduce.cpp

@@ -50,6 +50,9 @@ class DeductionWorklist {
   // Adds a single (param, arg) deduction of a specific.
   auto Add(SemIR::SpecificId param, SemIR::SpecificId arg,
            bool needs_substitution) -> void {
+    if (!param.is_valid() || !arg.is_valid()) {
+      return;
+    }
     auto& param_specific = context_.specifics().Get(param);
     auto& arg_specific = context_.specifics().Get(arg);
     if (param_specific.generic_id != arg_specific.generic_id) {

toolchain/check/handle_class.cpp

@@ -380,8 +380,8 @@ auto HandleParseNode(Context& context, Parse::AdaptDeclId node_id) -> bool {
     return true;
   }
 
-  auto adapted_type_id =
-      ExprAsType(context, node_id, adapted_type_expr_id).type_id;
+  auto [adapted_inst_id, adapted_type_id] =
+      ExprAsType(context, node_id, adapted_type_expr_id);
   adapted_type_id = context.AsCompleteType(
       adapted_type_id,
       [&] {
@@ -405,13 +405,13 @@ auto HandleParseNode(Context& context, Parse::AdaptDeclId node_id) -> bool {
 
   // Extend the class scope with the adapted type's scope if requested.
   if (introducer.modifier_set.HasAnyOf(KeywordModifierSet::Extend)) {
-    auto extended_scope_id = SemIR::NameScopeId::Invalid;
+    auto extended_scope_inst_id = SemIR::InstId::Invalid;
     if (adapted_type_id == SemIR::TypeId::Error) {
       // Recover by not extending any scope. We instead set has_error to true
       // below.
     } else if (auto* adapted_class_info =
                    TryGetAsClass(context, adapted_type_id)) {
-      extended_scope_id = adapted_class_info->scope_id;
+      extended_scope_inst_id = adapted_inst_id;
       CARBON_CHECK(adapted_class_info->scope_id.is_valid(),
                    "Complete class should have a scope");
     } else {
@@ -420,8 +420,8 @@ auto HandleParseNode(Context& context, Parse::AdaptDeclId node_id) -> bool {
     }
 
     auto& class_scope = context.name_scopes().Get(class_info.scope_id);
-    if (extended_scope_id.is_valid()) {
-      class_scope.extended_scopes.push_back(extended_scope_id);
+    if (extended_scope_inst_id.is_valid()) {
+      class_scope.extended_scopes.push_back(extended_scope_inst_id);
     } else {
       class_scope.has_error = true;
     }
@@ -448,9 +448,11 @@ struct BaseInfo {
 
   SemIR::TypeId type_id;
   SemIR::NameScopeId scope_id;
+  SemIR::InstId inst_id;
 };
 constexpr BaseInfo BaseInfo::Error = {.type_id = SemIR::TypeId::Error,
-                                      .scope_id = SemIR::NameScopeId::Invalid};
+                                      .scope_id = SemIR::NameScopeId::Invalid,
+                                      .inst_id = SemIR::InstId::Invalid};
 }  // namespace
 
 // Diagnoses an attempt to derive from a final type.
@@ -496,7 +498,9 @@ static auto CheckBaseType(Context& context, Parse::NodeId node_id,
 
   CARBON_CHECK(base_class_info->scope_id.is_valid(),
                "Complete class should have a scope");
-  return {.type_id = base_type_id, .scope_id = base_class_info->scope_id};
+  return {.type_id = base_type_id,
+          .scope_id = base_class_info->scope_id,
+          .inst_id = base_type_inst_id};
 }
 
 auto HandleParseNode(Context& context, Parse::BaseDeclId node_id) -> bool {
@@ -560,7 +564,7 @@ auto HandleParseNode(Context& context, Parse::BaseDeclId node_id) -> bool {
   if (introducer.modifier_set.HasAnyOf(KeywordModifierSet::Extend)) {
     auto& class_scope = context.name_scopes().Get(class_info.scope_id);
     if (base_info.scope_id.is_valid()) {
-      class_scope.extended_scopes.push_back(base_info.scope_id);
+      class_scope.extended_scopes.push_back(base_info.inst_id);
     } else {
       class_scope.has_error = true;
     }

toolchain/check/handle_impl.cpp

@@ -12,6 +12,7 @@
 #include "toolchain/check/modifiers.h"
 #include "toolchain/check/pattern_match.h"
 #include "toolchain/parse/typed_nodes.h"
+#include "toolchain/sem_ir/generic.h"
 #include "toolchain/sem_ir/ids.h"
 #include "toolchain/sem_ir/typed_insts.h"
 
@@ -173,30 +174,20 @@ static auto ExtendImpl(Context& context, Parse::NodeId extend_node,
     diag.Emit();
   }
 
-  auto facet_type = context.types().TryGetAs<SemIR::FacetType>(constraint_id);
-  if (!facet_type) {
+  if (!context.types().Is<SemIR::FacetType>(constraint_id)) {
     context.TODO(node_id, "extending non-facet-type constraint");
     parent_scope.has_error = true;
     return;
   }
-  const SemIR::FacetTypeInfo& info =
-      context.facet_types().Get(facet_type->facet_type_id);
-  for (auto interface_type : info.impls_constraints) {
-    auto& interface = context.interfaces().Get(interface_type.interface_id);
-    if (!interface.is_defined()) {
-      CARBON_DIAGNOSTIC(ExtendUndefinedInterface, Error,
-                        "`extend impl` requires a definition for interface {0}",
-                        InstIdAsType);
-      auto diag = context.emitter().Build(node_id, ExtendUndefinedInterface,
-                                          constraint_inst_id);
-      context.NoteUndefinedInterface(interface_type.interface_id, diag);
-      diag.Emit();
-      parent_scope.has_error = true;
-      return;
-    }
-
-    parent_scope.extended_scopes.push_back(interface.scope_id);
-  }
+  parent_scope.has_error |= !context.TryToDefineType(constraint_id, [&] {
+    CARBON_DIAGNOSTIC(ExtendUndefinedInterface, Error,
+                      "`extend impl` requires a definition for facet type {0}",
+                      InstIdAsType);
+    return context.emitter().Build(node_id, ExtendUndefinedInterface,
+                                   constraint_inst_id);
+  });
+
+  parent_scope.extended_scopes.push_back(constraint_inst_id);
 }
 
 // Pops the parameters of an `impl`, forming a `NameComponent` with no
@@ -333,6 +324,15 @@ static auto BuildImplDecl(Context& context, Parse::AnyImplDeclId node_id,
   // For an `extend impl` declaration, mark the impl as extending this `impl`.
   if (introducer.modifier_set.HasAnyOf(KeywordModifierSet::Extend)) {
     auto extend_node = introducer.modifier_node_id(ModifierOrder::Decl);
+    if (impl_info.generic_id.is_valid()) {
+      SemIR::TypeId type_id = context.insts().Get(constraint_inst_id).type_id();
+      constraint_inst_id = context.AddInst<SemIR::SpecificConstant>(
+          context.insts().GetLocId(constraint_inst_id),
+          {.type_id = type_id,
+           .inst_id = constraint_inst_id,
+           .specific_id =
+               context.generics().GetSelfSpecific(impl_info.generic_id)});
+    }
     ExtendImpl(context, extend_node, node_id, self_type_node, self_type_id,
                name.implicit_params_loc_id, constraint_inst_id,
                constraint_type_id);

toolchain/check/import_ref.cpp

@@ -1460,12 +1460,11 @@ class ImportRefResolver {
           context_.insts()
               .GetAs<SemIR::BaseDecl>(new_class.base_id)
               .base_type_id);
-      const auto& base_class = context_.classes().Get(
-          context_.insts().GetAs<SemIR::ClassType>(base_inst_id).class_id);
-      new_scope.extended_scopes.push_back(base_class.scope_id);
+      new_scope.extended_scopes.push_back(base_inst_id);
     }
-    CARBON_CHECK(new_scope.extended_scopes.size() ==
-                 import_scope.extended_scopes.size());
+    // TODO: `extended_scopes` from `extend impl` are currently not imported.
+    // CARBON_CHECK(new_scope.extended_scopes.size() ==
+    //              import_scope.extended_scopes.size());
   }
 
   auto TryResolveTypedInst(SemIR::ClassDecl inst,