AST: Add NormalProtocolConformance::AssociatedConformances

Author: slavapestov

include/swift/AST/ProtocolConformance.h

@@ -481,6 +481,7 @@ class NormalProtocolConformance : public RootProtocolConformance,
   /// Conformances that satisfy each of conformance requirements of the
   /// requirement signature of the protocol.
   ArrayRef<ProtocolConformanceRef> SignatureConformances;
+  MutableArrayRef<llvm::Optional<ProtocolConformanceRef>> AssociatedConformances;
 
   /// The lazy member loader provides callbacks for populating imported and
   /// deserialized conformances.
@@ -636,6 +637,16 @@ class NormalProtocolConformance : public RootProtocolConformance,
   ProtocolConformanceRef
   getAssociatedConformance(Type assocType, ProtocolDecl *protocol) const;
 
+  /// Allocate the backing array if needed, computing its size from the
+  ///protocol's requirement signature.
+  void createAssociatedConformanceArray();
+
+  llvm::Optional<ProtocolConformanceRef>
+  getAssociatedConformance(unsigned index) const;
+
+  void
+  setAssociatedConformance(unsigned index, ProtocolConformanceRef assocConf);
+
   /// Retrieve the value witness corresponding to the given requirement.
   Witness getWitness(ValueDecl *requirement) const;
 

lib/AST/ProtocolConformance.cpp

@@ -569,6 +569,42 @@ NormalProtocolConformance::getAssociatedConformance(Type assocType,
   return *result;
 }
 
+/// Allocate the backing array if needed, computing its size from the
+///protocol's requirement signature.
+void NormalProtocolConformance::createAssociatedConformanceArray() {
+  if (hasComputedAssociatedConformances())
+    return;
+
+  setHasComputedAssociatedConformances();
+
+  auto *proto = getProtocol();
+
+  unsigned count = 0;
+  for (auto req : proto->getRequirementSignature().getRequirements()) {
+    if (req.getKind() == RequirementKind::Conformance)
+      ++count;
+  }
+
+  auto &ctx = proto->getASTContext();
+  AssociatedConformances =
+      ctx.Allocate<llvm::Optional<ProtocolConformanceRef>>(count);
+}
+
+llvm::Optional<ProtocolConformanceRef>
+NormalProtocolConformance::getAssociatedConformance(unsigned index) const {
+  if (!hasComputedAssociatedConformances())
+    return llvm::None;
+
+  return AssociatedConformances[index];
+}
+
+void NormalProtocolConformance::setAssociatedConformance(
+    unsigned index, ProtocolConformanceRef assocConf) {
+  createAssociatedConformanceArray();
+
+  assert(!AssociatedConformances[index]);
+  AssociatedConformances[index] = assocConf;
+}
 
 /// A stripped-down version of Type::subst that only works on the protocol
 /// Self type wrapped in zero or more DependentMemberTypes.