AST: Add new variant of swift::getParameterAt() which takes pack expansions into account

Author: slavapestov

include/swift/AST/Decl.h

@@ -8399,6 +8399,10 @@ ParameterList *getParameterList(ValueDecl *source);
 /// there is none.
 ParameterList *getParameterList(DeclContext *source);
 
+/// Retrieve parameter declaration from the given source at given index, or
+/// nullptr if the source does not have a parameter list.
+const ParamDecl *getParameterAt(ConcreteDeclRef declRef, unsigned index);
+
 /// Retrieve parameter declaration from the given source at given index, or
 /// nullptr if the source does not have a parameter list.
 const ParamDecl *getParameterAt(const ValueDecl *source, unsigned index);

include/swift/AST/ParameterList.h

@@ -24,6 +24,8 @@
 
 namespace swift {
 
+class SubstitutionMap;
+
 /// This describes a list of parameters.  Each parameter descriptor is tail
 /// allocated onto this list.
 class alignas(ParamDecl *) ParameterList final :
@@ -134,6 +136,8 @@ class alignas(ParamDecl *) ParameterList final :
   /// based on the interface types of the parameters in this list.
   void getParams(SmallVectorImpl<AnyFunctionType::Param> &params) const;
 
+  unsigned getOrigParamIndex(SubstitutionMap subMap, unsigned substIndex) const;
+
   /// Return the full source range of this parameter.
   SourceRange getSourceRange() const;
   SourceLoc getStartLoc() const { return getSourceRange().Start; }

lib/AST/Decl.cpp

@@ -7743,6 +7743,17 @@ ParameterList *swift::getParameterList(DeclContext *source) {
   return nullptr;
 }
 
+const ParamDecl *swift::getParameterAt(ConcreteDeclRef declRef,
+                                       unsigned index) {
+  auto *source = declRef.getDecl();
+  if (auto *params = getParameterList(const_cast<ValueDecl *>(source))) {
+    unsigned origIndex = params->getOrigParamIndex(declRef.getSubstitutions(),
+                                                   index);
+    return params->get(origIndex);
+  }
+  return nullptr;
+}
+
 const ParamDecl *swift::getParameterAt(const ValueDecl *source,
                                        unsigned index) {
   if (auto *params = getParameterList(const_cast<ValueDecl *>(source))) {

lib/AST/ParameterPack.cpp

@@ -17,6 +17,7 @@
 
 #include "swift/AST/Types.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/ParameterList.h"
 #include "swift/AST/Type.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -249,4 +250,32 @@ PackType *PackType::flattenPackTypes() {
     return this;
 
   return PackType::get(getASTContext(), elts);
+}
+
+unsigned ParameterList::getOrigParamIndex(SubstitutionMap subMap,
+                                          unsigned substIndex) const {
+  unsigned remappedIndex = substIndex;
+
+  for (unsigned i = 0, e = size(); i < e; ++i) {
+    auto *param = get(i);
+    auto paramType = param->getInterfaceType();
+
+    unsigned substCount = 1;
+    if (auto *packExpansionType = paramType->getAs<PackExpansionType>()) {
+      auto replacementType = packExpansionType->getCountType().subst(subMap);
+      if (auto *packType = replacementType->getAs<PackType>()) {
+        substCount = packType->getNumElements();
+      }
+    }
+
+    if (remappedIndex < substCount)
+      return i;
+
+    remappedIndex -= substCount;
+  }
+
+  llvm::errs() << "Invalid substituted argument index: " << substIndex << "\n";
+  subMap.dump(llvm::errs());
+  dump(llvm::errs());
+  abort();
 }

lib/Sema/MiscDiagnostics.cpp

@@ -526,7 +526,7 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
         return;
 
       // Get underlying params for both callee and caller, if declared.
-      auto *calleeParam = getParameterAt(callee.getDecl(), argIndex);
+      auto *calleeParam = getParameterAt(callee, argIndex);
       auto *callerParam = dyn_cast_or_null<ParamDecl>(
           arg->getReferencedDecl(/*stopAtParenExpr=*/true).getDecl()
       );