[Sema] Copy key path component types when merging solutions

In 2eeff36 I forgot to copy key path component types when applying a solution to the constraint system. That caused a crash in key path code completion.

Fixes rdar://81118700 [SR-14979]

Author: ahoppen

include/swift/Sema/ConstraintSystem.h

@@ -1186,7 +1186,7 @@ class Solution {
   llvm::DenseMap<ASTNode, Type> nodeTypes;
 
   /// The key path component types introduced by this solution.
-  llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
+  llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, Type>
       keyPathComponentTypes;
 
   /// Contextual types introduced by this solution.
@@ -2183,9 +2183,23 @@ class ConstraintSystem {
   /// run through various diagnostics passes without actually mutating
   /// the types on the nodes.
   llvm::MapVector<ASTNode, Type> NodeTypes;
-  llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
+
+  /// The nodes for which we have produced types, along with the prior type
+  /// each node had before introducing this type.
+  llvm::SmallVector<std::pair<ASTNode, Type>, 8> addedNodeTypes;
+
+  /// Maps components in a key path expression to their type. Needed because
+  /// KeyPathExpr + Index isn't an \c ASTNode and thus can't be stored in \c
+  /// NodeTypes.
+  llvm::DenseMap<std::pair<const KeyPathExpr *, /*component index=*/unsigned>,
+                 Type>
       KeyPathComponentTypes;
 
+  /// Same as \c addedNodeTypes for \c KeyPathComponentTypes.
+  llvm::SmallVector<
+      std::tuple<const KeyPathExpr *, /*component index=*/unsigned, Type>>
+      addedKeyPathComponentTypes;
+
   /// Maps AST entries to their solution application targets.
   llvm::MapVector<SolutionApplicationTargetsKey, SolutionApplicationTarget>
     solutionApplicationTargets;
@@ -2266,10 +2280,6 @@ class ConstraintSystem {
   SmallVector<std::pair<ConstraintLocator *, OpenedArchetypeType *>, 4>
     OpenedExistentialTypes;
 
-  /// The nodes for which we have produced types, along with the prior type
-  /// each node had before introducing this type.
-  llvm::SmallVector<std::pair<ASTNode, Type>, 8> addedNodeTypes;
-
   /// The set of functions that have been transformed by a result builder.
   std::vector<std::pair<AnyFunctionRef, AppliedBuilderTransform>>
       resultBuilderTransformed;
@@ -2739,6 +2749,8 @@ class ConstraintSystem {
 
     unsigned numAddedNodeTypes;
 
+    unsigned numAddedKeyPathComponentTypes;
+
     unsigned numDisabledConstraints;
 
     unsigned numFavoredConstraints;
@@ -2954,10 +2966,15 @@ class ConstraintSystem {
   /// map is used throughout the expression type checker in order to
   /// avoid mutating expressions until we know we have successfully
   /// type-checked them.
-  void setType(KeyPathExpr *KP, unsigned I, Type T) {
+  void setType(const KeyPathExpr *KP, unsigned I, Type T) {
     assert(KP && "Expected non-null key path parameter!");
     assert(T && "Expected non-null type!");
-    KeyPathComponentTypes[std::make_pair(KP, I)] = T.getPointer();
+
+    Type &entry = KeyPathComponentTypes[{KP, I}];
+    Type oldType = entry;
+    entry = T;
+
+    addedKeyPathComponentTypes.push_back(std::make_tuple(KP, I, oldType));
   }
 
   /// Check to see if we have a type for a node.

lib/Sema/CSSolver.cpp

@@ -260,6 +260,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     setType(nodeType.first, nodeType.second);
   }
 
+  for (auto &nodeType : solution.keyPathComponentTypes) {
+    setType(nodeType.getFirst().first, nodeType.getFirst().second,
+            nodeType.getSecond());
+  }
+
   // Add the contextual types.
   for (const auto &contextualType : solution.contextualTypes) {
     if (!getContextualTypeInfo(contextualType.first)) {
@@ -489,6 +494,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
   numDefaultedConstraints = cs.DefaultedConstraints.size();
   numAddedNodeTypes = cs.addedNodeTypes.size();
+  numAddedKeyPathComponentTypes = cs.addedKeyPathComponentTypes.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
   numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
   numResultBuilderTransformed = cs.resultBuilderTransformed.size();
@@ -570,6 +576,19 @@ ConstraintSystem::SolverScope::~SolverScope() {
   }
   truncate(cs.addedNodeTypes, numAddedNodeTypes);
 
+  // Remove any node types we registered.
+  for (unsigned i : reverse(range(numAddedKeyPathComponentTypes,
+                                  cs.addedKeyPathComponentTypes.size()))) {
+    auto KeyPath = std::get<0>(cs.addedKeyPathComponentTypes[i]);
+    auto KeyPathIndex = std::get<1>(cs.addedKeyPathComponentTypes[i]);
+    if (Type oldType = std::get<2>(cs.addedKeyPathComponentTypes[i])) {
+      cs.KeyPathComponentTypes[{KeyPath, KeyPathIndex}] = oldType;
+    } else {
+      cs.KeyPathComponentTypes.erase({KeyPath, KeyPathIndex});
+    }
+  }
+  truncate(cs.addedKeyPathComponentTypes, numAddedKeyPathComponentTypes);
+
   /// Remove any builder transformed closures.
   truncate(cs.resultBuilderTransformed, numResultBuilderTransformed);
 

validation-test/IDE/crashers_fixed/sr14979.swift

@@ -0,0 +1,30 @@
+// RUN: %swift-ide-test --code-completion --source-filename %s -code-completion-token CC
+
+func foo() {
+    let months: [[String]] = []
+    let abcd: Int
+    let x = ForEach2(months) { group in
+        HStack2(spacing: 3) {
+            useKeyPath(id: \.abcd#^CC^#)
+        }
+    }
+}
+
+
+struct ForEach2<Data> where Data : RandomAccessCollection {
+    init(_ data: Data, content: @escaping (Data.Element) -> Void) {}
+}
+
+struct Bar {
+	let abcd: Int
+}
+
+func useKeyPath(id: KeyPath<Bar, String>) {}
+
+struct HStack2<Content> {
+    init(spacing: Double, @ViewBuilder2 content: () -> Content)
+}
+
+@resultBuilder struct ViewBuilder2 {
+  static func buildBlock<Content>(_ content: Content) -> Content { fatalError() }
+}