[AST] Always walk folded SequenceExpr if available

Expand the special-cased ASTWalker behavior for folded SequenceExprs
such that we always walk the folded expression when available. This
ensures that we don't attempt to add the same node multiple times
when expanding ASTScopes during pre-checking.

rdar://147751795

Author: hamishknight

include/swift/AST/ASTWalker.h

@@ -133,19 +133,6 @@ enum class QualifiedIdentTypeReprWalkingScheme {
   ASTOrderRecursive
 };
 
-/// Specifies the behavior for walking SequenceExprs.
-enum class SequenceWalking {
-  /// Walk into every element of the sequence, regardless of what state the
-  /// sequence is in.
-  Default,
-
-  /// If the sequence has been folded by type checking, only walk into the
-  /// elements that represent the operator nodes. This will ensure that the walk
-  /// does not visit the same AST nodes twice when it encounters a sequence that
-  /// has already been folded but hasn't been removed from the AST.
-  OnlyWalkFirstOperatorWhenFolded
-};
-
 /// An abstract class used to traverse an AST.
 class ASTWalker {
 public:
@@ -629,13 +616,6 @@ class ASTWalker {
     return MacroWalking::ArgumentsAndExpansion;
   }
 
-  /// This method configures how the walker should walk into SequenceExprs.
-  /// Needing to customize this behavior should be rare, as sequence expressions
-  /// are only encountered in un-typechecked ASTs.
-  virtual SequenceWalking getSequenceWalkingBehavior() const {
-    return SequenceWalking::Default;
-  }
-
   /// This method determines whether the given declaration should be
   /// considered to be in a macro expansion context. It can be configured
   /// by subclasses.

include/swift/AST/Expr.h

@@ -368,10 +368,9 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     IsObjC : 1
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(SequenceExpr, Expr, 32+1,
+  SWIFT_INLINE_BITFIELD_FULL(SequenceExpr, Expr, 32,
     : NumPadBits,
-    NumElements : 32,
-    IsFolded: 1
+    NumElements : 32
   );
 
   SWIFT_INLINE_BITFIELD(OpaqueValueExpr, Expr, 1,
@@ -3970,6 +3969,9 @@ class SequenceExpr final : public Expr,
     private llvm::TrailingObjects<SequenceExpr, Expr *> {
   friend TrailingObjects;
 
+  /// The cached folded expression.
+  Expr *FoldedExpr = nullptr;
+
   SequenceExpr(ArrayRef<Expr*> elements)
     : Expr(ExprKind::Sequence, /*Implicit=*/false) {
     Bits.SequenceExpr.NumElements = elements.size();
@@ -4005,11 +4007,14 @@ class SequenceExpr final : public Expr,
     getElements()[i] = e;
   }
 
-  bool isFolded() const {
-    return static_cast<bool>(Bits.SequenceExpr.IsFolded);
+  /// Retrieve the folded expression, or \c nullptr if the SequencExpr has
+  /// not yet been folded.
+  Expr *getFoldedExpr() const {
+    return FoldedExpr;
   }
-  void setFolded(bool folded) {
-    Bits.SequenceExpr.IsFolded = static_cast<unsigned>(folded);
+
+  void setFoldedExpr(Expr *foldedExpr) {
+    FoldedExpr = foldedExpr;
   }
 
   // Implement isa/cast/dyncast/etc.

lib/AST/ASTWalker.cpp

@@ -950,13 +950,18 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitSequenceExpr(SequenceExpr *E) {
-    if (Walker.getSequenceWalkingBehavior() ==
-            SequenceWalking::OnlyWalkFirstOperatorWhenFolded &&
-        E->isFolded()) {
-      if (E->getNumElements() > 1) {
-        if (Expr *Elt = doIt(E->getElement(1)))
-          E->setElement(1, Elt);
-      }
+    // After folding, the SequenceExpr elements can contain a broken mess of
+    // partially folded and/or duplicate nodes (since folding can mutate nodes
+    // in-place). We remove the SequenceExpr from the AST after folding, but
+    // there's still a period of time during pre-checking when it's still in the
+    // AST. To avoid issues for e.g ASTScope and availability scope
+    // construction, only walk the folded expression if we have it.
+    if (auto *foldedExpr = E->getFoldedExpr()) {
+      auto *newExpr = doIt(foldedExpr);
+      if (!newExpr)
+        return nullptr;
+
+      E->setFoldedExpr(newExpr);
       return E;
     }
 

lib/AST/AvailabilityScopeBuilder.cpp

@@ -176,15 +176,6 @@ class AvailabilityScopeBuilder : private ASTWalker {
     return MacroWalking::Arguments;
   }
 
-  SequenceWalking getSequenceWalkingBehavior() const override {
-    // Since availability scopes may be built at arbitrary times, the builder
-    // may encounter ASTs where SequenceExprs still exist and have not been
-    // folded, or it may encounter folded SequenceExprs that have not been
-    // removed from the AST. When folded exprs are encountered, its important
-    // to avoid walking into the same AST nodes twice.
-    return SequenceWalking::OnlyWalkFirstOperatorWhenFolded;
-  }
-
   /// Check whether this declaration is within a macro expansion buffer that
   /// will have its own availability scope that will be lazily expanded.
   bool isDeclInMacroExpansion(Decl *decl) const override {

lib/Sema/PreCheckTarget.cpp

@@ -1168,7 +1168,6 @@ class PreCheckTarget final : public ASTWalker {
       if (!result)
         return Action::Stop();
       // Already walked.
-      seqExpr->setFolded(true);
       return Action::SkipNode(result);
     }
 

lib/Sema/TypeCheckExpr.cpp

@@ -649,6 +649,7 @@ Expr *TypeChecker::foldSequence(SequenceExpr *expr, DeclContext *dc) {
   Expr *Result = ::foldSequence(dc, LHS, Elts, PrecedenceBound());
   assert(Elts.empty());
 
+  expr->setFoldedExpr(Result);
   return Result;
 }
 

test/NameLookup/scope_map_rdar147751795.swift

@@ -0,0 +1,19 @@
+// Note: test of the scope map. All of these tests are line- and
+// column-sensitive, so any additions should go at the end.
+
+struct X {}
+var y: () -> X
+
+struct S {
+  let x: () = y = {() -> X in
+    return .init()
+  }
+}
+
+// RUN: %target-swift-frontend -dump-scope-maps expanded %s 2>&1 | %FileCheck %s
+
+// rdar://147751795 - Make sure we don't end up with a duplicate closure scope.
+// CHECK:      PatternEntryInitializerScope {{.*}}, [8:15 - 10:3] entry 0 'x'
+// CHECK-NEXT:   `-ClosureParametersScope {{.*}}, [8:28 - 10:3]
+// CHECK-NEXT:     `-BraceStmtScope {{.*}}, [8:28 - 10:3]
+// CHECK-EMPTY: