AST: Introduce TypeBase::replaceDynamicSelfType()

Author: slavapestov

include/swift/AST/Types.h

@@ -1337,16 +1337,11 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// argument labels removed.
   Type removeArgumentLabels(unsigned numArgumentLabels);
 
-  /// Replace the base type of the result type of the given function
-  /// type with a new result type, as per a DynamicSelf or other
-  /// covariant return transformation.  The optionality of the
-  /// existing result will be preserved.
-  ///
-  /// \param newResultType The new result type.
-  ///
-  /// \param uncurryLevel The number of uncurry levels to apply before
-  /// replacing the type. With uncurry level == 0, this simply
-  /// replaces the current type with the new result type.
+  /// Replace DynamicSelfType anywhere it appears in covariant position with
+  /// the given type.
+  Type replaceDynamicSelfType(Type newSelfType);
+
+  /// Deprecated in favor of the above.
   Type replaceCovariantResultType(Type newResultType,
                                   unsigned uncurryLevel);
 

lib/AST/Type.cpp

@@ -1329,6 +1329,21 @@ Type TypeBase::removeArgumentLabels(unsigned numArgumentLabels) {
   return FunctionType::get(unlabeledParams, result, fnType->getExtInfo());
 }
 
+Type TypeBase::replaceDynamicSelfType(Type newSelfType) {
+  if (!hasDynamicSelfType())
+    return Type(this);
+
+  return Type(this).transformWithPosition(
+      TypePosition::Covariant,
+      [&](TypeBase *t, TypePosition pos) -> std::optional<Type> {
+      if (isa<DynamicSelfType>(t) &&
+          pos == TypePosition::Covariant) {
+        return newSelfType;
+      }
+    return std::nullopt;
+  });
+}
+
 Type TypeBase::replaceCovariantResultType(Type newResultType,
                                           unsigned uncurryLevel) {
   if (uncurryLevel == 0) {