Merge pull request #78706 from tshortli/sequence-expr-availability-scopes

AST: Fix SequenceExpr handling in AvailabilityScopeBuilder

Author: tshortli

include/swift/AST/ASTWalker.h

@@ -133,6 +133,19 @@ 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:
@@ -616,6 +629,13 @@ 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,9 +368,10 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     IsObjC : 1
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(SequenceExpr, Expr, 32,
+  SWIFT_INLINE_BITFIELD_FULL(SequenceExpr, Expr, 32+1,
     : NumPadBits,
-    NumElements : 32
+    NumElements : 32,
+    IsFolded: 1
   );
 
   SWIFT_INLINE_BITFIELD(OpaqueValueExpr, Expr, 1,
@@ -4002,6 +4003,13 @@ class SequenceExpr final : public Expr,
     getElements()[i] = e;
   }
 
+  bool isFolded() const {
+    return static_cast<bool>(Bits.SequenceExpr.IsFolded);
+  }
+  void setFolded(bool folded) {
+    Bits.SequenceExpr.IsFolded = static_cast<unsigned>(folded);
+  }
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Sequence;

lib/AST/ASTWalker.cpp

@@ -955,6 +955,16 @@ 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);
+      }
+      return E;
+    }
+
     for (unsigned i = 0, e = E->getNumElements(); i != e; ++i)
       if (Expr *Elt = doIt(E->getElement(i)))
         E->setElement(i, Elt);

lib/Sema/PreCheckTarget.cpp

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

lib/Sema/TypeCheckAvailability.cpp

@@ -538,6 +538,15 @@ 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 {

test/Sema/availability_scopes.swift

@@ -321,6 +321,22 @@ func testStringInterpolation() {
     """
 }
 
+// CHECK-NEXT: {{^}}  (decl_implicit version=50 decl=result
+// CHECK-NEXT: {{^}}    (decl_implicit version=50 decl=unusedA
+// CHECK-NEXT: {{^}}    (decl_implicit version=50 decl=unusedB
+
+func testSequenceExprs(b: Bool, x: Int?) {
+  let result = b
+    ? x.map {
+        let unusedA: Int
+        return $0
+      }
+    : x.map {
+        let unusedB: Int
+        return $0
+      }
+}
+
 // CHECK-NEXT: {{^}}  (decl version=50 unavailable=macOS decl=unavailableOnMacOS()
 // CHECK-NEXT: {{^}}    (decl_implicit version=50 unavailable=macOS decl=x