AST: Introduce ProtocolDecl::get{AssociatedType,ProtocolRequirement}()

Author: slavapestov

include/swift/AST/Decl.h

@@ -4195,6 +4195,14 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// a protocol having nested types (ObjC protocols).
   llvm::TinyPtrVector<AssociatedTypeDecl *> getAssociatedTypeMembers() const;
 
+  /// Returns a protocol requirement with the given name, or nullptr if the
+  /// name has multiple overloads, or no overloads at all.
+  ValueDecl *getSingleRequirement(DeclName name) const;
+
+  /// Returns an associated type with the given name, or nullptr if one does
+  /// not exist.
+  AssociatedTypeDecl *getAssociatedType(Identifier name) const;
+
   /// Walk this protocol and all of the protocols inherited by this protocol,
   /// transitively, invoking the callback function for each protocol.
   ///

lib/AST/ASTDemangler.cpp

@@ -84,11 +84,7 @@ TypeDecl *ASTBuilder::createTypeDecl(NodePointer node) {
       return nullptr;
 
     auto name = Ctx.getIdentifier(node->getChild(1)->getText());
-    auto results = proto->lookupDirect(name);
-    if (results.size() != 1)
-      return nullptr;
-
-    return dyn_cast<AssociatedTypeDecl>(results[0]);
+    return proto->getAssociatedType(name);
   }
 
   auto *DC = findDeclContext(node);
@@ -585,13 +581,8 @@ Type ASTBuilder::createDependentMemberType(StringRef member,
   if (!base->isTypeParameter())
     return Type();
 
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto member : protocol->lookupDirect(Ctx.getIdentifier(member),
-                                            flags)) {
-    if (auto assocType = dyn_cast<AssociatedTypeDecl>(member))
-      return DependentMemberType::get(base, assocType);
-  }
+  if (auto assocType = protocol->getAssociatedType(Ctx.getIdentifier(member)))
+    return DependentMemberType::get(base, assocType);
 
   return Type();
 }

lib/AST/Decl.cpp

@@ -4364,6 +4364,34 @@ ProtocolDecl::getAssociatedTypeMembers() const {
   return result;
 }
 
+ValueDecl *ProtocolDecl::getSingleRequirement(DeclName name) const {
+  auto results = const_cast<ProtocolDecl *>(this)->lookupDirect(name);
+  ValueDecl *result = nullptr;
+  for (auto candidate : results) {
+    if (candidate->getDeclContext() != this ||
+        !candidate->isProtocolRequirement())
+      continue;
+    if (result) {
+      // Multiple results.
+      return nullptr;
+    }
+    result = candidate;
+  }
+
+  return result;
+}
+
+AssociatedTypeDecl *ProtocolDecl::getAssociatedType(Identifier name) const {
+  auto results = const_cast<ProtocolDecl *>(this)->lookupDirect(name);
+  for (auto candidate : results) {
+    if (candidate->getDeclContext() == this &&
+        isa<AssociatedTypeDecl>(candidate)) {
+      return cast<AssociatedTypeDecl>(candidate);
+    }
+  }
+  return nullptr;
+}
+
 Type ProtocolDecl::getSuperclass() const {
   ASTContext &ctx = getASTContext();
   return evaluateOrDefault(ctx.evaluator,

lib/AST/ProtocolConformance.cpp

@@ -162,14 +162,7 @@ ProtocolConformanceRef::getTypeWitnessByName(Type type, Identifier name) const {
 
   // Find the named requirement.
   ProtocolDecl *proto = getRequirement();
-  AssociatedTypeDecl *assocType = nullptr;
-  auto members = proto->lookupDirect(name,
-                      NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions);
-  for (auto member : members) {
-    assocType = dyn_cast<AssociatedTypeDecl>(member);
-    if (assocType)
-      break;
-  }
+  auto *assocType = proto->getAssociatedType(name);
 
   // FIXME: Shouldn't this be a hard error?
   if (!assocType)
@@ -183,16 +176,7 @@ ConcreteDeclRef
 ProtocolConformanceRef::getWitnessByName(Type type, DeclName name) const {
   // Find the named requirement.
   auto *proto = getRequirement();
-  auto results =
-    proto->lookupDirect(name,
-                      NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions);
-
-  ValueDecl *requirement = nullptr;
-  for (auto *result : results) {
-    if (isa<ProtocolDecl>(result->getDeclContext()))
-      requirement = result;
-  }
-
+  auto *requirement = proto->getSingleRequirement(name);
   if (requirement == nullptr)
     return ConcreteDeclRef();
 

lib/SILGen/SILGen.cpp

@@ -219,16 +219,9 @@ FuncDecl *SILGenModule::getBridgeToObjectiveCRequirement(SILLocation loc) {
 
   // Look for _bridgeToObjectiveC().
   auto &ctx = getASTContext();
-  FuncDecl *found = nullptr;
   DeclName name(ctx, ctx.Id_bridgeToObjectiveC, llvm::ArrayRef<Identifier>());
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto member : proto->lookupDirect(name, flags)) {
-    if (auto func = dyn_cast<FuncDecl>(member)) {
-      found = func;
-      break;
-    }
-  }
+  auto *found = dyn_cast_or_null<FuncDecl>(
+    proto->getSingleRequirement(name));
 
   if (!found)
     diagnose(loc, diag::bridging_objcbridgeable_broken, name);
@@ -251,17 +244,10 @@ FuncDecl *SILGenModule::getUnconditionallyBridgeFromObjectiveCRequirement(
 
   // Look for _bridgeToObjectiveC().
   auto &ctx = getASTContext();
-  FuncDecl *found = nullptr;
   DeclName name(ctx, ctx.getIdentifier("_unconditionallyBridgeFromObjectiveC"),
                 llvm::makeArrayRef(Identifier()));
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto member : proto->lookupDirect(name, flags)) {
-    if (auto func = dyn_cast<FuncDecl>(member)) {
-      found = func;
-      break;
-    }
-  }
+  auto *found = dyn_cast_or_null<FuncDecl>(
+    proto->getSingleRequirement(name));
 
   if (!found)
     diagnose(loc, diag::bridging_objcbridgeable_broken, name);
@@ -284,19 +270,9 @@ SILGenModule::getBridgedObjectiveCTypeRequirement(SILLocation loc) {
 
   // Look for _bridgeToObjectiveC().
   auto &ctx = getASTContext();
-  AssociatedTypeDecl *found = nullptr;
-  DeclName name(ctx.Id_ObjectiveCType);
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto member : proto->lookupDirect(name, flags)) {
-    if (auto assocType = dyn_cast<AssociatedTypeDecl>(member)) {
-      found = assocType;
-      break;
-    }
-  }
-
+  auto *found = proto->getAssociatedType(ctx.Id_ObjectiveCType);
   if (!found)
-    diagnose(loc, diag::bridging_objcbridgeable_broken, name);
+    diagnose(loc, diag::bridging_objcbridgeable_broken, ctx.Id_ObjectiveCType);
 
   BridgedObjectiveCType = found;
   return found;
@@ -337,15 +313,8 @@ VarDecl *SILGenModule::getNSErrorRequirement(SILLocation loc) {
 
   // Look for _nsError.
   auto &ctx = getASTContext();
-  VarDecl *found = nullptr;
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto member : proto->lookupDirect(ctx.Id_nsError, flags)) {
-    if (auto var = dyn_cast<VarDecl>(member)) {
-      found = var;
-      break;
-    }
-  }
+  auto *found = dyn_cast_or_null<VarDecl>(
+      proto->getSingleRequirement(ctx.Id_nsError));
 
   NSErrorRequirement = found;
   return found;

lib/SILGen/SILGenExpr.cpp

@@ -3019,7 +3019,8 @@ getOrCreateKeyPathEqualsAndHash(SILGenModule &SGM,
     // Compare each pair of index values using the == witness from the
     // conformance.
     auto equatableProtocol = C.getProtocol(KnownProtocolKind::Equatable);
-    auto equalsMethod = equatableProtocol->lookupDirect(C.Id_EqualsOperator)[0];
+    auto equalsMethod = equatableProtocol->getSingleRequirement(
+      C.Id_EqualsOperator);
     auto equalsRef = SILDeclRef(equalsMethod);
     auto equalsTy = subSGF.SGM.Types.getConstantType(equalsRef);
 
@@ -3196,7 +3197,7 @@ getOrCreateKeyPathEqualsAndHash(SILGenModule &SGM,
       }
 
       VarDecl *hashValueVar =
-        cast<VarDecl>(hashableProto->lookupDirect(C.Id_hashValue)[0]);
+        cast<VarDecl>(hashableProto->getSingleRequirement(C.Id_hashValue));
 
       auto formalTy = index.FormalType;
       auto hashable = index.Hashable;

lib/SILOptimizer/IPO/EagerSpecializer.cpp

@@ -580,11 +580,10 @@ void EagerDispatch::emitRefCountedObjectCheck(SILBasicBlock *FailedTypeCheckBB,
                                                  {GenericMT});
   // Extract the i1 from the Bool struct.
   StructDecl *BoolStruct = cast<StructDecl>(Ctx.getBoolDecl());
-  auto Members = BoolStruct->lookupDirect(Ctx.Id_value_);
+  auto Members = BoolStruct->getStoredProperties();
   assert(Members.size() == 1 &&
          "Bool should have only one property with name '_value'");
-  auto Member = dyn_cast<VarDecl>(Members[0]);
-  assert(Member &&"Bool should have a property with name '_value' of type Int1");
+  auto Member = Members[0];
   auto BoolValue =
       Builder.emitStructExtract(Loc, IsClassRuntimeCheck, Member, BoolTy);
   Builder.createCondBranch(Loc, BoolValue, SuccessBB, FailedTypeCheckBB);

lib/SILOptimizer/Transforms/Outliner.cpp

@@ -142,18 +142,10 @@ static SILDeclRef getBridgeToObjectiveC(CanType NativeType,
     return SILDeclRef();
 
   auto Conformance = ConformanceRef->getConcrete();
-  FuncDecl *Requirement = nullptr;
   // bridgeToObjectiveC
   DeclName Name(Ctx, Ctx.Id_bridgeToObjectiveC, llvm::ArrayRef<Identifier>());
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto Member : Proto->lookupDirect(Name, flags)) {
-    if (auto Func = dyn_cast<FuncDecl>(Member)) {
-      Requirement = Func;
-      break;
-    }
-  }
-  assert(Requirement);
+  auto *Requirement = dyn_cast_or_null<FuncDecl>(
+    Proto->getSingleRequirement(Name));
   if (!Requirement)
     return SILDeclRef();
 
@@ -173,19 +165,11 @@ SILDeclRef getBridgeFromObjectiveC(CanType NativeType,
   if (!ConformanceRef)
     return SILDeclRef();
   auto Conformance = ConformanceRef->getConcrete();
-  FuncDecl *Requirement = nullptr;
   // _unconditionallyBridgeFromObjectiveC
   DeclName Name(Ctx, Ctx.getIdentifier("_unconditionallyBridgeFromObjectiveC"),
                 llvm::makeArrayRef(Identifier()));
-  auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-  flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-  for (auto Member : Proto->lookupDirect(Name, flags)) {
-    if (auto Func = dyn_cast<FuncDecl>(Member)) {
-      Requirement = Func;
-      break;
-    }
-  }
-  assert(Requirement);
+  auto *Requirement = dyn_cast_or_null<FuncDecl>(
+      Proto->getSingleRequirement(Name));
   if (!Requirement)
     return SILDeclRef();
 

lib/SILOptimizer/Utils/CastOptimizer.cpp

@@ -487,26 +487,6 @@ findBridgeToObjCFunc(SILOptFunctionBuilder &functionBuilder,
   (void)conf;
 
   // Generate code to invoke _bridgeToObjectiveC
-
-  auto *ntd = sourceType.getNominalOrBoundGenericNominal();
-  assert(ntd);
-  auto members = ntd->lookupDirect(mod.getASTContext().Id_bridgeToObjectiveC);
-  if (members.empty()) {
-    SmallVector<ValueDecl *, 4> foundMembers;
-    if (ntd->getDeclContext()->lookupQualified(
-            ntd, mod.getASTContext().Id_bridgeToObjectiveC,
-            NLOptions::NL_ProtocolMembers, foundMembers)) {
-      // Returned members are starting with the most specialized ones.
-      // Thus, the first element is what we are looking for.
-      members.push_back(foundMembers.front());
-    }
-  }
-
-  // There should be exactly one implementation of _bridgeToObjectiveC.
-  if (members.size() != 1)
-    return None;
-
-  auto bridgeFuncDecl = members.front();
   ModuleDecl *modDecl =
       mod.getASTContext().getLoadedModule(mod.getASTContext().Id_Foundation);
   if (!modDecl)
@@ -532,7 +512,7 @@ findBridgeToObjCFunc(SILOptFunctionBuilder &functionBuilder,
 
   // Get substitutions, if source is a bound generic type.
   auto subMap = sourceType->getContextSubstitutionMap(
-      mod.getSwiftModule(), bridgeFuncDecl->getDeclContext());
+      mod.getSwiftModule(), resultDecl->getDeclContext());
 
   // Implementation of _bridgeToObjectiveC could not be found.
   if (!bridgedFunc)

lib/Sema/CSApply.cpp

@@ -2127,16 +2127,15 @@ namespace {
         return witness;
       };
 
-      auto associatedTypeArray = 
+      auto *interpolationProto = 
         tc.getProtocol(expr->getLoc(),
-             KnownProtocolKind::ExpressibleByStringInterpolation)
-          ->lookupDirect(tc.Context.Id_StringInterpolation);
-      if (associatedTypeArray.empty()) {
+             KnownProtocolKind::ExpressibleByStringInterpolation);
+      auto associatedTypeDecl = interpolationProto->getAssociatedType(
+          tc.Context.Id_StringInterpolation);
+      if (associatedTypeDecl == nullptr) {
         tc.diagnose(expr->getStartLoc(), diag::interpolation_broken_proto);
         return nullptr;
       }
-      auto associatedTypeDecl =
-        cast<AssociatedTypeDecl>(associatedTypeArray.front());
       auto interpolationType =
         simplifyType(DependentMemberType::get(openedType, associatedTypeDecl));