Sema: Generalize analysis of same-type requirements of protocol extension witnesses

Fixes rdar://problem/122606610.

Author: slavapestov

lib/Sema/AssociatedTypeInference.cpp

@@ -1443,18 +1443,22 @@ static InferenceCandidateKind checkInferenceCandidate(
         }
 
         if (auto otherAssoc = other->getAs<DependentMemberType>()) {
-          if (otherAssoc->getBase()->isEqual(selfTy)) {
-            DependentMemberType *otherDMT;
-            if (ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
+          if (ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
 
-            otherDMT = otherAssoc;
-  
-            } else {
+          result->second = result->second.transform([&](Type t) -> Type{
+            if (t->isEqual(dmt))
+              return otherAssoc;
+            return t;
+          });
+          LLVM_DEBUG(llvm::dbgs() << "++ we can same-type to:\n";
+                     result->second->dump(llvm::dbgs()));
+          return InferenceCandidateKind::Good;
 
-            otherDMT = DependentMemberType::get(dmt->getBase(),
-                                                otherAssoc->getAssocType());
+          } else {
 
-            }
+          if (otherAssoc->getBase()->isEqual(selfTy)) {
+            auto *otherDMT = DependentMemberType::get(dmt->getBase(),
+                                                otherAssoc->getAssocType());
 
             result->second = result->second.transform([&](Type t) -> Type{
               if (t->isEqual(dmt))
@@ -1465,6 +1469,8 @@ static InferenceCandidateKind checkInferenceCandidate(
                        result->second->dump(llvm::dbgs()));
             return InferenceCandidateKind::Good;
           }
+
+          }
         }
         break;
       }

test/decl/protocol/req/rdar122596633-2.swift

@@ -0,0 +1,32 @@
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+
+public protocol P<A> {
+    associatedtype A
+    associatedtype B: P
+
+    func makeA() -> A
+    var b: B { get }
+}
+
+extension P where A == B.A {
+    public func makeA() -> B.A {
+        fatalError()
+    }
+}
+
+public struct S: P {
+    public var b: some P<Int> {
+        return G<Int>()
+    }
+}
+
+public struct G<A>: P {
+    public func makeA() -> A { fatalError() }
+    public var b: Never { fatalError() }
+}
+
+extension Never: P {
+    public func makeA() -> Never { fatalError() }
+    public var b: Never { fatalError() }
+}