[ConstraintSystem] Open only directly declared opaque types related to return statements (swiftlang#38762)

* [ConstraintSystem] NFC: Rename `openOpaqueTypeRec` and make `openOpaqueType` private for individual types

* [ConstraintSystem] Require a contextual purpose for `openOpaqueType(Type, ...)`

Some of the contexts require special handling e.g. return type of a function
can only open directly declared opaque result type(s), this would be
implemented in a follow-up commit.

* [ConstraintSystem] Open only directly declared opaque types related to return statements

Re-establish a check which would only open opaque types directly
declared in the return type of a function/accessor declaration,
otherwise constraint system would end up with type variables it
has no context for if the type had generic parameters.

Resolves: rdar://81531010

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -3855,11 +3855,15 @@ class ConstraintSystem {
   /// \returns The opened type, or \c type if there are no archetypes in it.
   Type openType(Type type, OpenedTypeMap &replacements);
 
-  /// "Open" an opaque archetype type.
-  Type openOpaqueType(Type type, ConstraintLocatorBuilder locator);
+private:
+  /// "Open" an opaque archetype type, similar to \c openType.
+  Type openOpaqueType(OpaqueTypeArchetypeType *type,
+                      ConstraintLocatorBuilder locator);
 
+public:
   /// Recurse over the given type and open any opaque archetype types.
-  Type openOpaqueTypeRec(Type type, ConstraintLocatorBuilder locator);
+  Type openOpaqueType(Type type, ContextualTypePurpose context,
+                      ConstraintLocatorBuilder locator);
 
   /// "Open" the given function type.
   ///

lib/Sema/BuilderTransform.cpp

@@ -1845,7 +1845,8 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
 
   // Bind the body result type to the type of the transformed expression.
   addConstraint(bodyResultConstraintKind, transformedType,
-                openOpaqueTypeRec(bodyResultType, locator), locator);
+                openOpaqueType(bodyResultType, CTP_ReturnStmt, locator),
+                locator);
   return getTypeMatchSuccess();
 }
 

lib/Sema/CSGen.cpp

@@ -1706,10 +1706,12 @@ namespace {
 
       auto locator = CS.getConstraintLocator(expr);
       auto contextualType = CS.getContextualType(expr);
+      auto contextualPurpose = CS.getContextualTypePurpose(expr);
 
       auto joinElementTypes = [&](Optional<Type> elementType) {
-        auto openedElementType = elementType.map(
-            [&](Type type) { return CS.openOpaqueTypeRec(type, locator); });
+        auto openedElementType = elementType.map([&](Type type) {
+          return CS.openOpaqueType(type, contextualPurpose, locator);
+        });
 
         const auto elements = expr->getElements();
         unsigned index = 0;
@@ -1829,7 +1831,9 @@ namespace {
 
       auto locator = CS.getConstraintLocator(expr);
       auto contextualType = CS.getContextualType(expr);
-      auto openedType = CS.openOpaqueTypeRec(contextualType, locator);
+      auto contextualPurpose = CS.getContextualTypePurpose(expr);
+      auto openedType =
+          CS.openOpaqueType(contextualType, contextualPurpose, locator);
 
       // If a contextual type exists for this expression and is a dictionary
       // type, apply it directly.
@@ -2241,7 +2245,8 @@ namespace {
         type = type->getReferenceStorageReferent();
 
         Type replacedType = CS.replaceInferableTypesWithTypeVars(type, locator);
-        Type openedType = CS.openOpaqueTypeRec(replacedType, locator);
+        Type openedType =
+            CS.openOpaqueType(replacedType, CTP_Initialization, locator);
         assert(openedType);
 
         auto *subPattern = cast<TypedPattern>(pattern)->getSubPattern();

lib/Sema/CSSimplify.cpp

@@ -12006,7 +12006,7 @@ void ConstraintSystem::addContextualConversionConstraint(
   // Add the constraint.
   auto *convertTypeLocator =
       getConstraintLocator(expr, LocatorPathElt::ContextualType(purpose));
-  auto openedType = openOpaqueTypeRec(conversionType, convertTypeLocator);
+  auto openedType = openOpaqueType(conversionType, purpose, convertTypeLocator);
   addConstraint(constraintKind, getType(expr), openedType, convertTypeLocator,
                 /*isFavored*/ true);
 }

lib/Sema/ConstraintSystem.cpp

@@ -844,9 +844,8 @@ Type ConstraintSystem::openType(Type type, OpenedTypeMap &replacements) {
     });
 }
 
-Type ConstraintSystem::openOpaqueType(Type type,
+Type ConstraintSystem::openOpaqueType(OpaqueTypeArchetypeType *opaque,
                                       ConstraintLocatorBuilder locator) {
-  auto opaque = type->castTo<OpaqueTypeArchetypeType>();
   auto opaqueLocator = locator.withPathElement(
       LocatorPathElt::OpenedOpaqueArchetype(opaque->getDecl()));
 
@@ -868,16 +867,36 @@ Type ConstraintSystem::openOpaqueType(Type type,
   return underlyingTyVar;
 }
 
-Type ConstraintSystem::openOpaqueTypeRec(Type type,
-                                         ConstraintLocatorBuilder locator) {
+Type ConstraintSystem::openOpaqueType(Type type, ContextualTypePurpose context,
+                                      ConstraintLocatorBuilder locator) {
   // Early return if `type` is `NULL` or if there are no opaque archetypes (in
   // which case there is certainly nothing for us to do).
   if (!type || !type->hasOpaqueArchetype())
     return type;
 
+  auto inReturnContext = [](ContextualTypePurpose context) {
+    return context == CTP_ReturnStmt || context == CTP_ReturnSingleExpr;
+  };
+
+  if (!(context == CTP_Initialization || inReturnContext(context)))
+    return type;
+
+  auto shouldOpen = [&](OpaqueTypeArchetypeType *opaqueType) {
+    if (!inReturnContext(context))
+      return true;
+
+    if (auto *func = dyn_cast<AbstractFunctionDecl>(DC))
+      return opaqueType->getDecl()->isOpaqueReturnTypeOfFunction(func);
+
+    return true;
+  };
+
   return type.transform([&](Type type) -> Type {
-    if (type->is<OpaqueTypeArchetypeType>())
-      return openOpaqueType(type, locator);
+    auto *opaqueType = type->getAs<OpaqueTypeArchetypeType>();
+
+    if (opaqueType && shouldOpen(opaqueType))
+      return openOpaqueType(opaqueType, locator);
+
     return type;
   });
 }

test/type/opaque_return_non_direct.swift

@@ -0,0 +1,39 @@
+// RUN: %target-swift-frontend -disable-availability-checking -typecheck -verify %s
+
+// rdar://81531010 - Failure to infer `Context` from contextual return type
+
+protocol View {}
+struct EmptyView: View {}
+
+@resultBuilder struct ViewBuilder {
+  static func buildBlock() -> EmptyView { EmptyView() }
+}
+
+struct Data: View {
+}
+
+protocol Modifier {
+  associatedtype Body
+  associatedtype Content
+
+  @ViewBuilder func body(content: Content) -> Body
+}
+
+extension View {
+  func config<S: Modifier>(_: S) -> S.Body where S.Content == Self, S.Body: View {
+    fatalError()
+  }
+}
+
+struct DataModifier<Content: View> : Modifier {
+  init(_: Data) {}
+  func body(content: Content) -> some View {}
+}
+
+struct Test {
+  typealias Value = DataModifier<Data>.Body
+
+  func test(data: Data) -> Value {
+    return data.config(DataModifier(data)) // Ok (Value is not considered an opaque type)
+  }
+}