Add LifetimeDependence to GenericFunctionType

Author: meg-gupta

include/swift/AST/Types.h

@@ -3866,10 +3866,11 @@ std::string getParamListAsString(ArrayRef<AnyFunctionType::Param> parameters);
 /// on those parameters and dependent member types thereof. The input and
 /// output types of the generic function can be expressed in terms of those
 /// generic parameters.
-class GenericFunctionType final : public AnyFunctionType,
-    public llvm::FoldingSetNode,
-    private llvm::TrailingObjects<GenericFunctionType, AnyFunctionType::Param,
-                                  Type> {
+class GenericFunctionType final
+    : public AnyFunctionType,
+      public llvm::FoldingSetNode,
+      private llvm::TrailingObjects<GenericFunctionType, AnyFunctionType::Param,
+                                    Type, LifetimeDependenceInfo> {
   friend TrailingObjects;
 
   GenericSignature Signature;
@@ -3882,6 +3883,10 @@ class GenericFunctionType final : public AnyFunctionType,
     return hasGlobalActor() + hasThrownError();
   }
 
+  size_t numTrailingObjects(OverloadToken<LifetimeDependenceInfo>) const {
+    return hasLifetimeDependenceInfo() ? 1 : 0;
+  }
+
   /// Construct a new generic function type.
   GenericFunctionType(GenericSignature sig, ArrayRef<Param> params, Type result,
                       llvm::Optional<ExtInfo> info, const ASTContext *ctx,
@@ -3909,7 +3914,17 @@ class GenericFunctionType final : public AnyFunctionType,
       return Type();
     return getTrailingObjects<Type>()[hasGlobalActor()];
   }
-                                    
+
+  LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+    if (!hasLifetimeDependenceInfo()) {
+      return LifetimeDependenceInfo();
+    }
+    auto *info = getTrailingObjects<LifetimeDependenceInfo>();
+    assert(!info->empty() && "If the LifetimeDependenceInfo was empty, we "
+                             "shouldn't have stored it.");
+    return *info;
+  }
+
   /// Retrieve the generic signature of this function type.
   GenericSignature getGenericSignature() const {
     return Signature;

lib/AST/ASTContext.cpp

@@ -4490,9 +4490,13 @@ GenericFunctionType *GenericFunctionType::get(GenericSignature sig,
   if (globalActor && !sig->isReducedType(globalActor))
     isCanonical = false;
 
+  bool hasLifetimeDependenceInfo =
+      info.has_value() && !info.value().getLifetimeDependenceInfo().empty();
+
   unsigned numTypes = (globalActor ? 1 : 0) + (thrownError ? 1 : 0);
-  size_t allocSize = totalSizeToAlloc<AnyFunctionType::Param, Type>(
-      params.size(), numTypes);
+  size_t allocSize =
+      totalSizeToAlloc<AnyFunctionType::Param, Type, LifetimeDependenceInfo>(
+          params.size(), numTypes, hasLifetimeDependenceInfo ? 1 : 0);
   void *mem = ctx.Allocate(allocSize, alignof(GenericFunctionType));
 
   auto properties = getGenericFunctionRecursiveProperties(params, result);
@@ -4522,6 +4526,11 @@ GenericFunctionType::GenericFunctionType(
     }
     if (Type thrownError = info->getThrownError())
       getTrailingObjects<Type>()[thrownErrorIndex] = thrownError;
+
+    auto lifetimeDependenceInfo = info->getLifetimeDependenceInfo();
+    if (!lifetimeDependenceInfo.empty()) {
+      *getTrailingObjects<LifetimeDependenceInfo>() = lifetimeDependenceInfo;
+    }
   }
 }
 

lib/AST/Type.cpp

@@ -4121,8 +4121,7 @@ LifetimeDependenceInfo AnyFunctionType::getLifetimeDependenceInfo() const {
   case TypeKind::Function:
     return cast<FunctionType>(this)->getLifetimeDependenceInfo();
   case TypeKind::GenericFunction:
-    // TODO: Handle GenericFunction
-    return LifetimeDependenceInfo();
+    return cast<GenericFunctionType>(this)->getLifetimeDependenceInfo();
 
   default:
     llvm_unreachable("Illegal type kind for AnyFunctionType.");