TypeDecoder: Implement lane-wise pack expansion

slavapestov · slavapestov · commit f40a44798033 · 2023-07-05T16:35:16.000-04:00
diff --git a/include/swift/AST/ASTDemangler.h b/include/swift/AST/ASTDemangler.h
@@ -64,12 +64,20 @@ class ASTBuilder {
   /// is a pack or not, so we have to find it here.
   GenericSignature GenericSig;
 
+  /// This builder doesn't perform "on the fly" substitutions, so we preserve
+  /// all pack expansions. We still need an active expansion stack though,
+  /// for the dummy implementation of these methods:
+  /// - beginPackExpansion()
+  /// - advancePackExpansion()
+  /// - createExpandedPackElement()
+  /// - endPackExpansion()
+  llvm::SmallVector<Type> ActivePackExpansions;
+
 public:
   using BuiltType = swift::Type;
   using BuiltTypeDecl = swift::GenericTypeDecl *; // nominal or type alias
   using BuiltProtocolDecl = swift::ProtocolDecl *;
   using BuiltGenericSignature = swift::GenericSignature;
-  using BuiltGenericTypeParam = swift::GenericTypeParamType *;
   using BuiltRequirement = swift::Requirement;
   using BuiltSubstitutionMap = swift::SubstitutionMap;
 
@@ -118,6 +126,14 @@ class ASTBuilder {
 
   Type createPackExpansionType(Type patternType, Type countType);
 
+  size_t beginPackExpansion(Type countType);
+
+  void advancePackExpansion(size_t index);
+
+  Type createExpandedPackElement(Type patternType);
+
+  void endPackExpansion();
+
   Type createFunctionType(
       ArrayRef<Demangle::FunctionParam<Type>> params,
       Type output, FunctionTypeFlags flags,
diff --git a/include/swift/Demangling/TypeDecoder.h b/include/swift/Demangling/TypeDecoder.h
@@ -478,7 +478,6 @@ class TypeDecoder {
   using BuiltSubstitution = typename BuilderType::BuiltSubstitution;
   using BuiltRequirement = typename BuilderType::BuiltRequirement;
   using BuiltLayoutConstraint = typename BuilderType::BuiltLayoutConstraint;
-  using BuiltGenericTypeParam = typename BuilderType::BuiltGenericTypeParam;
   using BuiltGenericSignature = typename BuilderType::BuiltGenericSignature;
   using BuiltSubstitutionMap = typename BuilderType::BuiltSubstitutionMap;
   using NodeKind = Demangle::Node::Kind;
@@ -886,10 +885,11 @@ class TypeDecoder {
 
       bool hasParamFlags = false;
       llvm::SmallVector<FunctionParam<BuiltType>, 8> parameters;
-      if (!decodeMangledFunctionInputType(Node->getChild(firstChildIdx),
-                                          depth + 1, parameters, hasParamFlags))
-        return MAKE_NODE_TYPE_ERROR0(Node->getChild(firstChildIdx),
-                                     "failed to decode function type");
+      auto optError = decodeMangledFunctionInputType(Node->getChild(firstChildIdx),
+                                                     depth + 1, parameters, hasParamFlags);
+      if (optError)
+        return *optError;
+
       flags =
           flags.withNumParameters(parameters.size())
               .withParameterFlags(hasParamFlags)
@@ -1034,24 +1034,31 @@ class TypeDecoder {
           return MAKE_NODE_TYPE_ERROR0(element->getChild(typeChildIndex),
                                        "no children");
         }
+
+        StringRef label;
         if (element->getChild(typeChildIndex)->getKind() == NodeKind::TupleElementName) {
-          labels.push_back(element->getChild(typeChildIndex)->getText());
+          label = element->getChild(typeChildIndex)->getText();
           typeChildIndex++;
-
-        // Otherwise, add an empty label.
-        } else {
-          labels.push_back(StringRef());
         }
 
         // Decode the element type.
-        auto elementType =
-            decodeMangledType(element->getChild(typeChildIndex), depth + 1,
-                              /*forRequirement=*/false);
-        if (elementType.isError())
-          return elementType;
-
-        elements.push_back(elementType.getType());
+        auto optError = decodeTypeSequenceElement(
+            element->getChild(typeChildIndex), depth + 1,
+            [&](BuiltType type) {
+              elements.push_back(type);
+              labels.push_back(label);
+            });
+        if (optError)
+          return *optError;
       }
+
+      // Unwrap one-element tuples.
+      //
+      // FIXME: The behavior of one-element labeled tuples is inconsistent throughout
+      // the different re-implementations of type substitution and pack expansion.
+      // if (elements.size() == 1)
+      //  return elements[0];
+
       return Builder.createTupleType(elements, labels);
     }
     case NodeKind::TupleElement:
@@ -1077,12 +1084,13 @@ class TypeDecoder {
 
       for (auto &element : *Node) {
         // Decode the element type.
-        auto elementType =
-            decodeMangledType(element, depth + 1, /*forRequirement=*/false);
-        if (elementType.isError())
-          return elementType;
-
-        elements.push_back(elementType.getType());
+        auto optError = decodeTypeSequenceElement(
+            element, depth + 1,
+            [&](BuiltType elementType) {
+              elements.push_back(elementType);
+            });
+        if (optError)
+          return *optError;
       }
 
       switch (Node->getKind()) {
@@ -1259,9 +1267,9 @@ return {}; // Not Implemented!
                                            /*forRequirement=*/false);
           if (substTy.isError())
             return substTy;
-          substitutions.emplace_back(
-              Builder.createGenericTypeParameterType(paramDepth, index),
-              substTy.getType());
+          auto paramTy = Builder.createGenericTypeParameterType(
+              paramDepth, index);
+          substitutions.emplace_back(paramTy, substTy.getType());
           ++index;
         }
       }
@@ -1364,6 +1372,58 @@ return {}; // Not Implemented!
   }
 
 private:
+  template<typename Fn>
+  llvm::Optional<TypeLookupError>
+  decodeTypeSequenceElement(Demangle::NodePointer node, unsigned depth,
+                            Fn resultCallback) {
+    if (node->getKind() == NodeKind::Type)
+      node = node->getChild(0);
+
+    if (node->getKind() == NodeKind::PackExpansion) {
+      if (node->getNumChildren() < 2)
+        return MAKE_NODE_TYPE_ERROR(node,
+                                    "fewer children (%zu) than required (2)",
+                                    node->getNumChildren());
+
+      auto patternType = node->getChild(0);
+
+      // Decode the shape pack first, to form a metadata pack.
+      auto countType = decodeMangledType(node->getChild(1), depth);
+      if (countType.isError())
+        return *countType.getError();
+
+      // Push the pack expansion on the active expansion stack inside the
+      // builder concept.
+      size_t numElements = Builder.beginPackExpansion(countType.getType());
+
+      for (size_t i = 0; i < numElements; ++i) {
+        // Advance the lane index inside the builder concept.
+        Builder.advancePackExpansion(i);
+
+        // Decode the pattern type, taking the ith element of each pack
+        // referenced therein.
+        auto expandedElementType = decodeMangledType(patternType, depth);
+        if (expandedElementType.isError())
+          return *expandedElementType.getError();
+
+        resultCallback(Builder.createExpandedPackElement(
+            expandedElementType.getType()));
+      }
+
+      // Pop the active expansion stack inside the builder concept.
+      Builder.endPackExpansion();
+    } else {
+      auto elementType =
+          decodeMangledType(node, depth, /*forRequirement=*/false);
+      if (elementType.isError())
+        return *elementType.getError();
+
+      resultCallback(elementType.getType());
+    }
+
+    return llvm::None;
+  }
+
   template <typename T>
   bool decodeImplFunctionPart(Demangle::NodePointer node, unsigned depth,
                               llvm::SmallVectorImpl<T> &results) {
@@ -1527,12 +1587,12 @@ return {}; // Not Implemented!
     return Builder.createProtocolDecl(node);
   }
 
-  bool decodeMangledFunctionInputType(
+  llvm::Optional<TypeLookupError> decodeMangledFunctionInputType(
       Demangle::NodePointer node, unsigned depth,
       llvm::SmallVectorImpl<FunctionParam<BuiltType>> &params,
       bool &hasParamFlags) {
     if (depth > TypeDecoder::MaxDepth)
-      return false;
+      return llvm::None;
 
     // Look through a couple of sugar nodes.
     if (node->getKind() == NodeKind::Type ||
@@ -1543,7 +1603,7 @@ return {}; // Not Implemented!
 
     auto decodeParamTypeAndFlags =
         [&](Demangle::NodePointer typeNode,
-            FunctionParam<BuiltType> &param) -> bool {
+            FunctionParam<BuiltType> &param) -> llvm::Optional<TypeLookupError> {
       Demangle::NodePointer node = typeNode;
 
       bool recurse = true;
@@ -1592,17 +1652,15 @@ return {}; // Not Implemented!
         }
       }
 
-      auto paramType = decodeMangledType(node, depth + 1,
-                                         /*forRequirement=*/false);
-      if (paramType.isError())
-        return false;
-
-      param.setType(paramType.getType());
-      return true;
+      return decodeTypeSequenceElement(node, depth + 1,
+                                       [&](BuiltType paramType) {
+                                         param.setType(paramType);
+                                         params.push_back(param);
+                                       });
     };
 
     auto decodeParam =
-        [&](NodePointer paramNode) -> llvm::Optional<FunctionParam<BuiltType>> {
+        [&](NodePointer paramNode) -> llvm::Optional<TypeLookupError> {
       if (paramNode->getKind() != NodeKind::TupleElement)
         return llvm::None;
 
@@ -1618,40 +1676,41 @@ return {}; // Not Implemented!
           hasParamFlags = true;
           break;
 
-        case NodeKind::Type:
-          if (!decodeParamTypeAndFlags(child->getFirstChild(), param))
-            return llvm::None;
+        case NodeKind::Type: {
+          auto optError = decodeParamTypeAndFlags(
+              child->getFirstChild(), param);
+          if (optError)
+            return optError;
           break;
+        }
 
         default:
-          return llvm::None;
+          return TYPE_LOOKUP_ERROR_FMT("unknown node");
         }
       }
 
-      return param;
+      return llvm::None;
     };
 
     // Expand a single level of tuple.
     if (node->getKind() == NodeKind::Tuple) {
       // Decode all the elements as separate arguments.
       for (const auto &elt : *node) {
-        auto param = decodeParam(elt);
-        if (!param)
-          return false;
-
-        params.push_back(std::move(*param));
+        auto optError = decodeParam(elt);
+        if (optError)
+          return *optError;
       }
 
-      return true;
+      return llvm::None;
     }
 
     // Otherwise, handle the type as a single argument.
     FunctionParam<BuiltType> param;
-    if (!decodeParamTypeAndFlags(node, param))
-      return false;
+    auto optError = decodeParamTypeAndFlags(node, param);
+    if (optError)
+      return *optError;
 
-    params.push_back(std::move(param));
-    return true;
+    return llvm::None;
   }
 };
 
diff --git a/include/swift/Remote/MetadataReader.h b/include/swift/Remote/MetadataReader.h
@@ -179,7 +179,6 @@ class MetadataReader {
   using BuiltRequirement = typename BuilderType::BuiltRequirement;
   using BuiltSubstitution = typename BuilderType::BuiltSubstitution;
   using BuiltSubstitutionMap = typename BuilderType::BuiltSubstitutionMap;
-  using BuiltGenericTypeParam = typename BuilderType::BuiltGenericTypeParam;
   using BuiltGenericSignature = typename BuilderType::BuiltGenericSignature;
   using StoredPointer = typename Runtime::StoredPointer;
   using StoredSignedPointer = typename Runtime::StoredSignedPointer;
diff --git a/include/swift/RemoteInspection/TypeRefBuilder.h b/include/swift/RemoteInspection/TypeRefBuilder.h
@@ -764,6 +764,24 @@ class TypeRefBuilder {
     return nullptr;
   }
 
+  size_t beginPackExpansion(const TypeRef *countType) {
+    // FIXME: Remote mirrors support for variadic generics.
+    return 0;
+  }
+
+  void advancePackExpansion(size_t index) {
+    // FIXME: Remote mirrors support for variadic generics.
+  }
+
+  const TypeRef *createExpandedPackElement(const TypeRef *patternType) {
+    // FIXME: Remote mirrors support for variadic generics.
+    return nullptr;
+  }
+
+  void endPackExpansion() {
+    // FIXME: Remote mirrors support for variadic generics.
+  }
+
   const FunctionTypeRef *createFunctionType(
       llvm::ArrayRef<remote::FunctionParam<const TypeRef *>> params,
       const TypeRef *result, FunctionTypeFlags flags,
diff --git a/lib/AST/ASTDemangler.cpp b/lib/AST/ASTDemangler.cpp
@@ -359,6 +359,26 @@ Type ASTBuilder::createPackExpansionType(Type patternType, Type countType) {
   return PackExpansionType::get(patternType, countType);
 }
 
+size_t ASTBuilder::beginPackExpansion(Type countType) {
+  ActivePackExpansions.push_back(countType);
+
+  return 1;
+}
+
+void ASTBuilder::advancePackExpansion(size_t index) {
+  assert(index == 0);
+}
+
+Type ASTBuilder::createExpandedPackElement(Type patternType) {
+  assert(!ActivePackExpansions.empty());
+  auto countType = ActivePackExpansions.back();
+  return PackExpansionType::get(patternType, countType);
+}
+
+void ASTBuilder::endPackExpansion() {
+  ActivePackExpansions.pop_back();
+}
+
 Type ASTBuilder::createFunctionType(
     ArrayRef<Demangle::FunctionParam<Type>> params,
     Type output, FunctionTypeFlags flags,
@@ -876,7 +896,7 @@ Type ASTBuilder::createParenType(Type base) {
 GenericSignature
 ASTBuilder::createGenericSignature(ArrayRef<BuiltType> builtParams,
                                    ArrayRef<BuiltRequirement> requirements) {
-  std::vector<BuiltGenericTypeParam> params;
+  std::vector<GenericTypeParamType *> params;
   for (auto &param : builtParams) {
     auto paramTy = param->getAs<GenericTypeParamType>();
     if (!paramTy)
diff --git a/stdlib/public/runtime/MetadataLookup.cpp b/stdlib/public/runtime/MetadataLookup.cpp
