[CS] Remove isInputExpression parameter

This wasn't consistently used, and consequently
could result in some expressions getting their
parents invalidated. Instead, replace it with a
query to make sure we don't try and add an
expression we've already computed the parent info
for.

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -4562,8 +4562,7 @@ class ConstraintSystem {
   ///
   /// \returns a possibly-sanitized expression, or null if an error occurred.
   [[nodiscard]]
-  Expr *generateConstraints(Expr *E, DeclContext *dc,
-                            bool isInputExpression = true);
+  Expr *generateConstraints(Expr *E, DeclContext *dc);
 
   /// Generate constraints for binding the given pattern to the
   /// value of the given expression.

lib/Sema/CSGen.cpp

@@ -4959,10 +4959,8 @@ bool ConstraintSystem::generateConstraints(
   }
 }
 
-Expr *ConstraintSystem::generateConstraints(
-    Expr *expr, DeclContext *dc, bool isInputExpression) {
-  if (isInputExpression)
-    InputExprs.insert(expr);
+Expr *ConstraintSystem::generateConstraints(Expr *expr, DeclContext *dc) {
+  InputExprs.insert(expr);
   return generateConstraintsFor(*this, expr, dc);
 }
 

lib/Sema/CSSyntacticElement.cpp

@@ -1115,8 +1115,8 @@ class SyntacticElementConstraintGenerator
 
       for (auto node : braceStmt->getElements()) {
         if (auto expr = node.dyn_cast<Expr *>()) {
-          auto generatedExpr = cs.generateConstraints(
-            expr, context.getAsDeclContext(), /*isInputExpression=*/false);
+          auto generatedExpr =
+              cs.generateConstraints(expr, context.getAsDeclContext());
           if (!generatedExpr) {
             hadError = true;
           }
@@ -1250,8 +1250,7 @@ class SyntacticElementConstraintGenerator
       auto *expr = returnStmt->getResult();
       assert(expr && "single expression closure without expression?");
 
-      expr = cs.generateConstraints(expr, context.getAsDeclContext(),
-                                    /*isInputExpression=*/false);
+      expr = cs.generateConstraints(expr, context.getAsDeclContext());
       if (!expr) {
         hadError = true;
         return;

lib/Sema/ConstraintSystem.cpp

@@ -822,6 +822,12 @@ void ConstraintSystem::addPackEnvironment(PackElementExpr *packElement,
 static void extendDepthMap(
    Expr *expr,
    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap) {
+  // If we already have an entry in the map, we don't need to update it. This
+  // avoids invalidating previous entries when solving a smaller component of a
+  // larger AST node, e.g during conjunction solving.
+  if (depthMap.contains(expr))
+    return;
+
   class RecordingTraversal : public ASTWalker {
     SmallVector<ClosureExpr *, 4> Closures;
 

unittests/Sema/ConstraintGenerationTests.cpp

@@ -38,7 +38,7 @@ TEST_F(SemaTest, TestImplicitForceCastConstraintGeneration) {
   auto *castExpr = ForcedCheckedCastExpr::createImplicit(Context, literal,
                                                          Context.TheAnyType);
 
-  auto *expr = cs.generateConstraints(castExpr, DC, /*isInputExpression=*/true);
+  auto *expr = cs.generateConstraints(castExpr, DC);
 
   ASSERT_NE(expr, nullptr);
 
@@ -66,7 +66,7 @@ TEST_F(SemaTest, TestImplicitCoercionConstraintGeneration) {
   auto *castExpr = CoerceExpr::createImplicit(Context, literal,
                                               getStdlibType("Double"));
 
-  auto *expr = cs.generateConstraints(castExpr, DC, /*isInputExpression=*/true);
+  auto *expr = cs.generateConstraints(castExpr, DC);
 
   ASSERT_NE(expr, nullptr);
 
@@ -95,7 +95,7 @@ TEST_F(SemaTest, TestImplicitConditionalCastConstraintGeneration) {
   auto *castExpr = ConditionalCheckedCastExpr::createImplicit(
       Context, literal, getStdlibType("Double"));
 
-  auto *expr = cs.generateConstraints(castExpr, DC, /*isInputExpression=*/true);
+  auto *expr = cs.generateConstraints(castExpr, DC);
 
   ASSERT_NE(expr, nullptr);
 
@@ -179,3 +179,44 @@ TEST_F(SemaTest, TestCaptureListIsNotOpenedEarly) {
     ASSERT_TRUE(cs.hasType(capture.getVar()));
   }
 }
+
+TEST_F(SemaTest, TestMultiStmtClosureBodyParentAndDepth) {
+  // {
+  //   ()
+  //   return ()
+  // }
+  DeclAttributes attrs;
+  auto *closure = new (Context) ClosureExpr(attrs,
+                                            /*braceRange=*/SourceRange(),
+                                            /*capturedSelfDecl=*/nullptr,
+                                            ParameterList::createEmpty(Context),
+                                            /*asyncLoc=*/SourceLoc(),
+                                            /*throwsLoc=*/SourceLoc(),
+                                            /*thrownType*/ nullptr,
+                                            /*arrowLoc=*/SourceLoc(),
+                                            /*inLoc=*/SourceLoc(),
+                                            /*explicitResultType=*/nullptr, DC);
+  closure->setImplicit();
+
+  auto *RS = ReturnStmt::createImplicit(
+      Context, TupleExpr::createImplicit(Context, {}, {}));
+
+  closure->setBody(BraceStmt::createImplicit(Context, {
+    TupleExpr::createImplicit(Context, {}, {}), RS
+  }));
+
+  SyntacticElementTarget target(closure, DC, ContextualTypeInfo(),
+                                /*isDiscarded*/ true);
+
+  ConstraintSystem cs(DC, ConstraintSystemOptions());
+  cs.solve(target);
+
+  ASSERT_EQ(cs.getParentExpr(closure), nullptr);
+  ASSERT_EQ(cs.getExprDepth(closure), 0);
+
+  // We visit the ReturnStmt twice when computing the parent map, ensure we
+  // don't invalidate its parent on the second walk during the conjunction.
+  auto *result = RS->getResult();
+  ASSERT_EQ(cs.getParentExpr(result), closure);
+  ASSERT_EQ(cs.getExprDepth(result), 1);
+}