[AST] Make 'StmtConditionElement' a single 'PointerUnion'

Previously, 'IntroducerLoc' and 'ThePattern' were only used for pattern
binidng cases. Create a new 'ConditionalPatternBindingInfo' type to
cover such cases, and make 'StmtConditionElement' a pure 'PointerUnion'
type.

This makes it clear which fields are used in which condition kind. Also,
we can expect overall size reduction of StmtCondition when the
majority of the conditions are simple boolean expressions.

Author: rintaro

include/swift/AST/Stmt.h

@@ -386,7 +386,47 @@ class DeferStmt : public Stmt {
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Defer; }
 };
 
-  
+/// Represent `let`/`var` optional binding, or `case` pattern matching in
+/// conditional statements (i.e. `if`, `guard`, `while`).
+class alignas(8) ConditionalPatternBindingInfo
+    : public ASTAllocated<ConditionalPatternBindingInfo> {
+public:
+  /// Location of the var/let/case keyword.
+  SourceLoc IntroducerLoc;
+
+  /// Pattern being matched. In the case of an "implicit optional" pattern, the
+  /// OptionalSome pattern is explicitly added to this as an 'implicit' pattern.
+  Pattern *ThePattern;
+
+  /// The value for matching.
+  Expr *Initializer;
+
+  ConditionalPatternBindingInfo(SourceLoc IntroducerLoc, Pattern *ThePattern,
+                                Expr *Initializer)
+      : IntroducerLoc(IntroducerLoc), ThePattern(ThePattern),
+        Initializer(Initializer) {}
+
+public:
+  static ConditionalPatternBindingInfo *create(ASTContext &ctx,
+                                               SourceLoc IntroducerLoc,
+                                               Pattern *ThePattern,
+                                               Expr *Initializer) {
+    return new (ctx)
+        ConditionalPatternBindingInfo(IntroducerLoc, ThePattern, Initializer);
+  }
+
+  SourceLoc getIntroducerLoc() const { return IntroducerLoc; }
+  void setIntroducerLoc(SourceLoc Loc) { IntroducerLoc = Loc; }
+  Pattern *getPattern() const { return ThePattern; }
+  void setPattern(Pattern *P) { ThePattern = P; }
+  Expr *getInitializer() const { return Initializer; }
+  void setInitializer(Expr *E) { Initializer = E; }
+
+  SourceRange getSourceRange() const;
+  SourceLoc getStartLoc() const { return getSourceRange().Start; };
+  SourceLoc getEndLoc() const { return getSourceRange().End; };
+};
+
 /// An expression that guards execution based on whether the run-time
 /// configuration supports a given API, e.g.,
 /// #available(OSX >= 10.9, iOS >= 7.0).
@@ -505,46 +545,39 @@ class PoundHasSymbolInfo final : public ASTAllocated<PoundHasSymbolInfo> {
   }
 };
 
-/// This represents an entry in an "if" or "while" condition.  Pattern bindings
-/// can bind any number of names in the pattern binding decl, and may have an
-/// associated where clause.  When "if let" is involved, an arbitrary number of
-/// pattern bindings and conditional expressions are permitted, e.g.:
+/// This represents an entry in an "if" or "while" condition.
+/// Either a boolean expression, optional binding, pattern matching,
+/// `#available`, or `#_hasSymbol`.
 ///
-///   if let x = ..., y = ... where x > y,
-///      let z = ...
-/// which would be represented as four StmtConditionElement entries, one for
-/// the "x" binding, one for the "y" binding, one for the where clause, one for
-/// "z"'s binding.  A simple "if" statement is represented as a single binding.
+/// E.g. this 'if' statement has 5 'StmtConditionElement'.
+///   if
+///     list.count == 1,                          // CK_Boolean
+///     let firstElem = list.first,               // CK_PatternBinding
+///     case .foo(let value?, "int") = firstElem, // CK_PatternBinding
+///     #available(myOS 13),                      // CK_Availability
+///     #_hasSymbol(MyStruct.peform(operation:))  // CK_HasSymbol
+///   { ... }
 ///
 class alignas(1 << PatternAlignInBits) StmtConditionElement {
-  /// If this is a pattern binding, it may be the first one in a declaration, in
-  /// which case this is the location of the var/let/case keyword.  If this is
-  /// the second pattern (e.g. for 'y' in "var x = ..., y = ...") then this
-  /// location is invalid.
-  SourceLoc IntroducerLoc;
-
-  /// In a pattern binding, this is pattern being matched.  In the case of an
-  /// "implicit optional" pattern, the OptionalSome pattern is explicitly added
-  /// to this as an 'implicit' pattern.
-  Pattern *ThePattern = nullptr;
-
-  /// This is either the boolean condition, the #available information, or
-  /// the #_hasSymbol information.
-  llvm::PointerUnion<Expr *, PoundAvailableInfo *, PoundHasSymbolInfo *>
+private:
+  llvm::PointerUnion<Expr *, ConditionalPatternBindingInfo *,
+                     PoundAvailableInfo *, PoundHasSymbolInfo *>
       Condition;
 
 public:
-  StmtConditionElement() {}
-  StmtConditionElement(SourceLoc IntroducerLoc, Pattern *ThePattern, Expr *Init)
-      : IntroducerLoc(IntroducerLoc), ThePattern(ThePattern), Condition(Init) {}
+  StmtConditionElement() : Condition(nullptr) {}
   StmtConditionElement(Expr *cond) : Condition(cond) {}
-
+  StmtConditionElement(ConditionalPatternBindingInfo *Info) : Condition(Info) {}
   StmtConditionElement(PoundAvailableInfo *Info) : Condition(Info) {}
-
   StmtConditionElement(PoundHasSymbolInfo *Info) : Condition(Info) {}
 
-  SourceLoc getIntroducerLoc() const { return IntroducerLoc; }
-  void setIntroducerLoc(SourceLoc loc) { IntroducerLoc = loc; }
+  static StmtConditionElement fromOpaqueValue(void *opaque) {
+    StmtConditionElement val;
+    val.Condition = decltype(Condition)::getFromOpaqueValue(opaque);
+    return val;
+  }
+
+  void *getOpaqueValue() const { return Condition.getOpaqueValue(); }
 
   /// ConditionKind - This indicates the sort of condition this is.
   enum ConditionKind {
@@ -555,20 +588,19 @@ class alignas(1 << PatternAlignInBits) StmtConditionElement {
   };
 
   ConditionKind getKind() const {
-    if (ThePattern)
-      return CK_PatternBinding;
     if (Condition.is<Expr *>())
       return CK_Boolean;
+    if (Condition.is<ConditionalPatternBindingInfo *>())
+      return CK_PatternBinding;
     if (Condition.is<PoundAvailableInfo *>())
       return CK_Availability;
-    assert(Condition.is<PoundHasSymbolInfo *>());
-    return CK_HasSymbol;
+    if (Condition.is<PoundHasSymbolInfo *>())
+      return CK_HasSymbol;
+    return CK_Boolean;
   }
 
   /// Boolean Condition Accessors.
-  Expr *getBooleanOrNull() const {
-    return getKind() == CK_Boolean ? Condition.get<Expr *>() : nullptr;
-  }
+  Expr *getBooleanOrNull() const { return Condition.dyn_cast<Expr *>(); }
 
   Expr *getBoolean() const {
     assert(getKind() == CK_Boolean && "Not a condition");
@@ -580,35 +612,43 @@ class alignas(1 << PatternAlignInBits) StmtConditionElement {
   }
 
   /// Pattern Binding Accessors.
-  Pattern *getPatternOrNull() const {
-    return ThePattern;
+  ConditionalPatternBindingInfo *getPatternBindingOrNull() const {
+    return Condition.dyn_cast<ConditionalPatternBindingInfo *>();
   }
 
-  Pattern *getPattern() const {
+  ConditionalPatternBindingInfo *getPatternBinding() const {
     assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
-    return ThePattern;
+    return Condition.get<ConditionalPatternBindingInfo *>();
   }
 
-  void setPattern(Pattern *P) {
-    assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
-    ThePattern = P;
+  SourceLoc getIntroducerLoc() const {
+    return getPatternBinding()->getIntroducerLoc();
   }
-  
-  /// Pattern Binding Accessors.
-  Expr *getInitializerOrNull() const {
-    return getKind() == CK_PatternBinding ? Condition.get<Expr *>() : nullptr;
+
+  void setIntroducerLoc(SourceLoc loc) {
+    getPatternBinding()->setIntroducerLoc(loc);
   }
 
-  Expr *getInitializer() const {
-    assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
-    return Condition.get<Expr *>();
+  Pattern *getPatternOrNull() const {
+    if (auto *binding = getPatternBindingOrNull())
+      return binding->getPattern();
+    return nullptr;
   }
-  
-  void setInitializer(Expr *E) {
-    assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
-    Condition = E;
+
+  Pattern *getPattern() const { return getPatternBinding()->getPattern(); }
+
+  void setPattern(Pattern *P) { getPatternBinding()->setPattern(P); }
+
+  Expr *getInitializerOrNull() const {
+    if (auto *binding = getPatternBindingOrNull())
+      return binding->getInitializer();
+    return nullptr;
   }
-  
+
+  Expr *getInitializer() const { return getPatternBinding()->getInitializer(); }
+
+  void setInitializer(Expr *E) { getPatternBinding()->setInitializer(E); }
+
   // Availability Accessors
   PoundAvailableInfo *getAvailability() const {
     assert(getKind() == CK_Availability && "Not an #available condition");

lib/AST/Stmt.cpp

@@ -571,6 +571,21 @@ bool StmtConditionElement::rebindsSelf(ASTContext &Ctx,
   return false;
 }
 
+SourceRange ConditionalPatternBindingInfo::getSourceRange() const {
+  SourceLoc Start;
+  if (IntroducerLoc.isValid())
+    Start = IntroducerLoc;
+  else
+    Start = ThePattern->getStartLoc();
+
+  SourceLoc End = Initializer->getEndLoc();
+  if (Start.isValid() && End.isValid()) {
+    return SourceRange(Start, End);
+  } else {
+    return SourceRange();
+  }
+}
+
 PoundAvailableInfo *
 PoundAvailableInfo::create(ASTContext &ctx, SourceLoc PoundLoc,
                            SourceLoc LParenLoc,
@@ -603,18 +618,7 @@ SourceRange StmtConditionElement::getSourceRange() const {
   case StmtConditionElement::CK_HasSymbol:
     return getHasSymbolInfo()->getSourceRange();
   case StmtConditionElement::CK_PatternBinding:
-    SourceLoc Start;
-    if (IntroducerLoc.isValid())
-      Start = IntroducerLoc;
-    else
-      Start = getPattern()->getStartLoc();
-    
-    SourceLoc End = getInitializer()->getEndLoc();
-    if (Start.isValid() && End.isValid()) {
-      return SourceRange(Start, End);
-    } else {
-      return SourceRange();
-    }
+    return getPatternBinding()->getSourceRange();
   }
 
   llvm_unreachable("Unhandled StmtConditionElement in switch.");
@@ -636,7 +640,7 @@ SourceLoc StmtConditionElement::getStartLoc() const {
   case StmtConditionElement::CK_Availability:
     return getAvailability()->getStartLoc();
   case StmtConditionElement::CK_PatternBinding:
-    return getSourceRange().Start;
+    return getPatternBinding()->getStartLoc();
   case StmtConditionElement::CK_HasSymbol:
     return getHasSymbolInfo()->getStartLoc();
   }
@@ -651,7 +655,7 @@ SourceLoc StmtConditionElement::getEndLoc() const {
   case StmtConditionElement::CK_Availability:
     return getAvailability()->getEndLoc();
   case StmtConditionElement::CK_PatternBinding:
-    return getSourceRange().End;
+    return getPatternBinding()->getEndLoc();
   case StmtConditionElement::CK_HasSymbol:
     return getHasSymbolInfo()->getEndLoc();
   }

lib/Parse/ParseStmt.cpp

@@ -1862,7 +1862,8 @@ Parser::parseStmtConditionElement(SmallVectorImpl<StmtConditionElement> &result,
                                 ErrorExpr(ThePattern.get()->getEndLoc()));
   }
 
-  result.push_back({IntroducerLoc, ThePattern.get(), Init.get()});
+  result.push_back(ConditionalPatternBindingInfo::create(
+      Context, IntroducerLoc, ThePattern.get(), Init.get()));
   return Status;
 }
 

lib/Sema/DerivedConformanceCodable.cpp

@@ -1676,8 +1676,8 @@ deriveBodyDecodable_enum_init(AbstractFunctionDecl *initDecl, void *) {
           C, VarDecl::Introducer::Let,
           NamedPattern::createImplicit(C, theKeyDecl));
 
-      guardElements.emplace_back(SourceLoc(), theKeyPattern,
-                                 allKeysPopFirstCallExpr);
+      guardElements.emplace_back(ConditionalPatternBindingInfo::create(
+          C, SourceLoc(), theKeyPattern, allKeysPopFirstCallExpr));
 
       // generate: allKeys.isEmpty;
       auto *allKeysIsEmptyExpr =

lib/Sema/TypeCheckStorage.cpp

@@ -1644,7 +1644,8 @@ synthesizeLazyGetterBody(AccessorDecl *Get, VarDecl *VD, VarDecl *Storage,
     createPropertyLoadOrCallSuperclassGetter(Get, Storage,
                                              TargetImpl::Storage, Ctx);
   SmallVector<StmtConditionElement, 1> Cond;
-  Cond.emplace_back(SourceLoc(), Some, StoredValueExpr);
+  Cond.emplace_back(ConditionalPatternBindingInfo::create(
+      Ctx, SourceLoc(), Some, StoredValueExpr));
 
   // Build the early return inside the if.
   auto *Tmp1DRE = new (Ctx) DeclRefExpr(Tmp1VD, DeclNameLoc(), /*Implicit*/true,

unittests/AST/SourceLocTests.cpp

@@ -162,36 +162,39 @@ TEST(SourceLoc, StmtConditionElement) {
                                             , false);
 
   // Case a, when the IntroducerLoc is valid.
-  auto introducer = StmtConditionElement( start.getAdvancedLoc(3)
-                                        , pattern, init);
-  
+  auto introducer = StmtConditionElement(ConditionalPatternBindingInfo::create(
+      C.Ctx, start.getAdvancedLoc(3), pattern, init));
+
   EXPECT_EQ(start.getAdvancedLoc(3), introducer.getStartLoc());
   EXPECT_EQ(start.getAdvancedLoc(25), introducer.getEndLoc());
   EXPECT_EQ( SourceRange(start.getAdvancedLoc(3), start.getAdvancedLoc(25))
            , introducer.getSourceRange());
 
   // Case b, when the IntroducerLoc is invalid, but the pattern has a valid loc.
-  auto patternStmtCond = StmtConditionElement(SourceLoc(), pattern, init);
-  
+  auto patternStmtCond = StmtConditionElement(
+      ConditionalPatternBindingInfo::create(C.Ctx, SourceLoc(), pattern, init));
+
   EXPECT_EQ(start.getAdvancedLoc(7), patternStmtCond.getStartLoc());
   EXPECT_EQ(start.getAdvancedLoc(25), patternStmtCond.getEndLoc());
   EXPECT_EQ( SourceRange(start.getAdvancedLoc(7), start.getAdvancedLoc(25))
            , patternStmtCond.getSourceRange());
 
   // If the IntroducerLoc is valid but the stmt cond init is invalid.
   auto invalidInit = new (C.Ctx) IntegerLiteralExpr("1", SourceLoc(), false);
-  auto introducerStmtInvalid = StmtConditionElement( start.getAdvancedLoc(3)
-                                                   , pattern, invalidInit);
-  
+  auto introducerStmtInvalid =
+      StmtConditionElement(ConditionalPatternBindingInfo::create(
+          C.Ctx, start.getAdvancedLoc(3), pattern, invalidInit));
+
   EXPECT_EQ(SourceLoc(), introducerStmtInvalid.getStartLoc());
   EXPECT_EQ(SourceLoc(), introducerStmtInvalid.getEndLoc());
   EXPECT_EQ(SourceRange(), introducerStmtInvalid.getSourceRange());
 
   // If the IntroducerLoc is invalid, the pattern is valid, but the stmt cond 
   // init is invalid.
-  auto patternStmtInvalid = StmtConditionElement( SourceLoc(), pattern
-                                                , invalidInit);
-  
+  auto patternStmtInvalid =
+      StmtConditionElement(ConditionalPatternBindingInfo::create(
+          C.Ctx, SourceLoc(), pattern, invalidInit));
+
   EXPECT_EQ(SourceLoc(), patternStmtInvalid.getStartLoc());
   EXPECT_EQ(SourceLoc(), patternStmtInvalid.getEndLoc());
   EXPECT_EQ(SourceRange(), patternStmtInvalid.getSourceRange());