Implement arity reabstraction for closures

include/swift/AST/Types.h

@@ -3103,6 +3103,8 @@ class AnyFunctionType : public TypeBase {
         return true;
       return false;
     }
+
+    Param getCanonical(CanGenericSignature genericSig) const;
 
     ParameterTypeFlags getParameterFlags() const { return Flags; }
 
@@ -6864,6 +6866,8 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
 BEGIN_CAN_TYPE_WRAPPER(PackType, Type)
   static CanPackType get(const ASTContext &ctx, ArrayRef<CanType> elements);
   static CanPackType get(const ASTContext &ctx, CanTupleEltTypeArrayRef elts);
+  static CanPackType get(const ASTContext &ctx,
+                         AnyFunctionType::CanParamArrayRef params);
 
   static CanTypeWrapper<PackType>
   getSingletonPackExpansion(CanType packParameter);

include/swift/Basic/ArrayRefView.h

@@ -167,6 +167,17 @@ template<typename Orig, typename Projected>
 using CastArrayRefView =
   ArrayRefView<Orig, Projected *, arrayRefViewCastHelper<Projected, Orig>>;
 
+namespace generator_details {
+template <class T> struct is_array_ref_like;
+
+template <class Orig, class Projected, Projected (&Project)(const Orig &),
+          bool AllowOrigAccess>
+struct is_array_ref_like<ArrayRefView<Orig, Projected, Project,
+                                      AllowOrigAccess>> {
+  enum { value = true };
+};
+}
+
 } // end namespace swift
 
 #endif // SWIFT_BASIC_ARRAYREFVIEW_H

include/swift/Basic/Generators.h

@@ -0,0 +1,213 @@
+//===--- Generators.h - "Coroutines" for doing traversals -------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a few types for defining types that follow this
+// simple generator concept:
+//
+// concept Generator {
+//   // ...some number of accessors for the current value...
+//
+//   /// Is this generator finished producing values?
+//   bool isFinished() const;
+//
+//   /// Given that this generator is not finished, advance to the
+//   /// next value.
+//   void advance();
+//
+//   /// Finish the generator, asserting that all values have been
+//   /// produced.
+//   void finish();
+// };
+//
+// concept SimpleGenerator : Generator {
+//   type reference;
+//
+//   reference claimNext();
+// }
+//
+// Generators are useful when some structure needs to be traversed but
+// that traversal can't be done in a simple lexical loop.  For example,
+// you can't do two traversals in parallel with a single loop unless you
+// break down all the details of the traversal.  This is a minor problem
+// for simple traversals like walking a flat array, but it's a significant
+// problem when traversals get more complex, like when different components
+// of an array are grouped together according to some additional structure
+// (such as the abstraction pattern of a function's parameter list).
+// It's tempting to write those traversals as higher-order functions that
+// invoke a callback for each element, but this breaks down when parallel
+// traversal is required.  Expressing the traversal as a generator
+// allows the traversal logic to to be reused without that limitation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_GENERATORS_H
+#define SWIFT_BASIC_GENERATORS_H
+
+#include "llvm/ADT/ArrayRef.h"
+
+namespace swift {
+
+namespace generator_details {
+
+template <class T>
+struct is_array_ref_like {
+  enum { value = false };
+};
+
+template <class T>
+struct is_array_ref_like<llvm::ArrayRef<T>> {
+  enum { value = true };
+};
+
+template <class T>
+struct is_array_ref_like<llvm::MutableArrayRef<T>> {
+  enum { value = true };
+};
+}
+
+/// A class for generating the elements of an ArrayRef-like collection.
+template <class CollectionType>
+class ArrayRefGenerator {
+  static_assert(generator_details::is_array_ref_like<CollectionType>::value,
+                "ArrayRefGenerator should only be used with ArrayRef-like "
+                "types");
+
+  CollectionType values;
+
+public:
+  using reference =
+    typename std::iterator_traits<typename CollectionType::iterator>::reference;
+
+  ArrayRefGenerator() {}
+  ArrayRefGenerator(CollectionType values) : values(values) {}
+
+  // Prevent accidental copying of the generator.
+  ArrayRefGenerator(const ArrayRefGenerator &other) = delete;
+  ArrayRefGenerator &operator=(const ArrayRefGenerator &other) = delete;
+
+  ArrayRefGenerator(ArrayRefGenerator &&other) = default;
+  ArrayRefGenerator &operator=(ArrayRefGenerator &&other) = default;
+
+  /// Explicitly copy the current generator state.
+  ArrayRefGenerator clone() const {
+    return ArrayRefGenerator(values);
+  }
+
+  /// Return the current element of the array.
+  reference getCurrent() const {
+    assert(!isFinished());
+    return values.front();
+  }
+
+  /// Claim the current element of the array and advance past it.
+  reference claimNext() {
+    assert(!isFinished());
+    reference result = getCurrent();
+    advance();
+    return result;
+  }
+
+  /// Claim the next N elements of the array and advance past them.
+  CollectionType claimNext(size_t count) {
+    assert(count <= values.size() && "claiming too many values");
+    CollectionType result = values.slice(0, count);
+    values = values.slice(count);
+    return result;
+  }
+
+  /// Is this generation finished?
+  bool isFinished() const {
+    return values.empty();
+  }
+
+  /// Given that this generation is not finished, advance to the
+  /// next element.
+  void advance() {
+    assert(!isFinished());
+    values = values.slice(1);
+  }
+
+  /// Perform any final work required to complete the generation.
+  void finish() {
+    assert(isFinished() && "didn't finish generating the collection");
+  }
+};
+
+/// An abstracting reference to an existing generator.
+///
+/// The implementation of this type holds the reference to the existing
+/// generator without allocating any additional storage; it is sufficient
+/// for the caller ensures that the object passed to the constructor
+/// stays valid.  Values of this type can otherwise be safely copied
+/// around.
+template <class T>
+class SimpleGeneratorRef {
+public:
+  using reference = T;
+
+private:
+  struct VTable {
+    bool (*isFinished)(const void *impl);
+    reference (*claimNext)(void *impl);
+    void (*advance)(void *impl);
+    void (*finish)(void *impl);
+  };
+
+  template <class G> struct VTableImpl {
+    static constexpr VTable vtable = {
+      [](const void *p) { return static_cast<const G*>(p)->isFinished(); },
+      [](void *p) -> reference { return static_cast<G*>(p)->claimNext(); },
+      [](void *p) { static_cast<G*>(p)->advance(); },
+      [](void *p) { static_cast<G*>(p)->finish(); },
+    };
+  };
+
+  const VTable *vtable;
+  void *pointer;
+
+public:
+  constexpr SimpleGeneratorRef() : vtable(nullptr), pointer(nullptr) {}
+
+  template <class G>
+  constexpr SimpleGeneratorRef(G &generator)
+    : vtable(&VTableImpl<G>::vtable), pointer(&generator) {}
+
+  /// Test whether this generator ref was initialized with a
+  /// valid reference to a generator.
+  explicit operator bool() const {
+    return pointer != nullptr;
+  }
+
+  bool isFinished() const {
+    assert(pointer);
+    return vtable->isFinished(pointer);
+  }
+
+  reference claimNext() {
+    assert(pointer);
+    return vtable->claimNext(pointer);
+  }
+
+  void advance() {
+    assert(pointer);
+    vtable->advance(pointer);
+  }
+
+  void finish() {
+    assert(pointer);
+    vtable->finish(pointer);
+  }
+};
+
+} // end namespace swift
+
+#endif

include/swift/SIL/SILArgument.h

@@ -342,12 +342,14 @@ class SILFunctionArgument : public SILArgument {
       ValueOwnershipKind ownershipKind, const ValueDecl *decl = nullptr,
       bool isNoImplicitCopy = false,
       LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None,
-      bool isCapture = false)
+      bool isCapture = false,
+      bool isParameterPack = false)
       : SILArgument(ValueKind::SILFunctionArgument, parentBlock, type,
                     ownershipKind, decl) {
     sharedUInt32().SILFunctionArgument.noImplicitCopy = isNoImplicitCopy;
     sharedUInt32().SILFunctionArgument.lifetimeAnnotation = lifetimeAnnotation;
     sharedUInt32().SILFunctionArgument.closureCapture = isCapture;
+    sharedUInt32().SILFunctionArgument.parameterPack = isParameterPack;
   }
 
   // A special constructor, only intended for use in
@@ -373,6 +375,23 @@ class SILFunctionArgument : public SILArgument {
     sharedUInt32().SILFunctionArgument.closureCapture = newValue;
   }
 
+  /// Is this parameter a pack that corresponds to multiple
+  /// formal parameters?  (This could mean multiple ParamDecl*s,
+  /// or it could mean a ParamDecl* that's a pack expansion.)  Note
+  /// that not all lowered parameters of pack type are parameter packs:
+  /// they can be part of a single formal parameter of tuple type.
+  /// This flag indicates that the lowered parameter has a one-to-many
+  /// relationship with formal parameters.
+  ///
+  /// TODO: preserve the parameter pack references in SIL in a side table
+  /// instead of using a single bit.
+  bool isFormalParameterPack() const {
+    return sharedUInt32().SILFunctionArgument.parameterPack;
+  }
+  void setFormalParameterPack(bool isPack) {
+    sharedUInt32().SILFunctionArgument.parameterPack = isPack;
+  }
+
   LifetimeAnnotation getLifetimeAnnotation() const {
     return LifetimeAnnotation::Case(
         sharedUInt32().SILFunctionArgument.lifetimeAnnotation);
@@ -416,6 +435,7 @@ class SILFunctionArgument : public SILArgument {
     setNoImplicitCopy(arg->isNoImplicitCopy());
     setLifetimeAnnotation(arg->getLifetimeAnnotation());
     setClosureCapture(arg->isClosureCapture());
+    setFormalParameterPack(arg->isFormalParameterPack());
   }
 
   static bool classof(const SILInstruction *) = delete;

include/swift/SIL/SILNode.h

@@ -268,7 +268,8 @@ class alignas(8) SILNode :
     SHARED_FIELD(StringLiteralInst, uint32_t length);
     SHARED_FIELD(PointerToAddressInst, uint32_t alignment);
     SHARED_FIELD(SILFunctionArgument, uint32_t noImplicitCopy : 1,
-                 lifetimeAnnotation : 2, closureCapture : 1);
+                 lifetimeAnnotation : 2, closureCapture : 1,
+                 parameterPack : 1);
 
     // Do not use `_sharedUInt32_private` outside of SILNode.
   } _sharedUInt32_private;

lib/AST/ASTContext.cpp

@@ -3305,6 +3305,16 @@ CanPackType CanPackType::get(const ASTContext &C,
   return CanPackType(PackType::get(C, ncElements));
 }
 
+CanPackType CanPackType::get(const ASTContext &C,
+                             AnyFunctionType::CanParamArrayRef params) {
+  SmallVector<Type, 8> ncElements;
+  ncElements.reserve(params.size());
+  for (auto param : params) {
+    ncElements.push_back(param.getParameterType());
+  }
+  return CanPackType(PackType::get(C, ncElements));
+}
+
 PackType *PackType::get(const ASTContext &C, ArrayRef<Type> elements) {
   RecursiveTypeProperties properties;
   bool isCanonical = true;

lib/AST/Type.cpp

@@ -1606,14 +1606,19 @@ getCanonicalParams(AnyFunctionType *funcType,
                    SmallVectorImpl<AnyFunctionType::Param> &canParams) {
   auto origParams = funcType->getParams();
   for (auto param : origParams) {
-    // Canonicalize the type and drop the internal label to canonicalize the
-    // Param.
-    canParams.emplace_back(param.getPlainType()->getReducedType(genericSig),
-                           param.getLabel(), param.getParameterFlags(),
-                           /*InternalLabel=*/Identifier());
+    canParams.emplace_back(param.getCanonical(genericSig));
   }
 }
 
+AnyFunctionType::Param
+AnyFunctionType::Param::getCanonical(CanGenericSignature genericSig) const {
+  // Canonicalize the type and drop the internal label to canonicalize the
+  // Param.
+  return Param(getPlainType()->getReducedType(genericSig),
+               getLabel(), getParameterFlags(),
+               /*InternalLabel=*/Identifier());
+}
+
 CanType TypeBase::computeCanonicalType() {
   assert(!hasCanonicalTypeComputed() && "called unnecessarily");
 

lib/SIL/IR/AbstractionPattern.cpp

@@ -2310,10 +2310,13 @@ const {
   if (yieldType)
     yieldType = yieldType->getReducedType(substSig);
 
+  // Note that we specifically do not want to put subMap in the
+  // abstraction patterns here, because the types we will be lowering
+  // against them will not be substituted.
   return std::make_tuple(
-          AbstractionPattern(subMap, substSig, substTy->getReducedType(substSig)),
+          AbstractionPattern(substSig, substTy->getReducedType(substSig)),
           subMap,
           yieldType
-            ? AbstractionPattern(subMap, substSig, yieldType)
+            ? AbstractionPattern(substSig, yieldType)
             : AbstractionPattern::getInvalid());
 }