AST: Store captured PackElementExprs in CaptureInfo

Author: slavapestov

include/swift/AST/CaptureInfo.h

@@ -41,9 +41,12 @@ template <> struct DenseMapInfo<swift::CapturedValue>;
 namespace swift {
 class ValueDecl;
 class FuncDecl;
+class Expr;
 class OpaqueValueExpr;
+class PackElementExpr;
 class VarDecl;
 class GenericEnvironment;
+class Type;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -52,7 +55,7 @@ class CapturedValue {
 
 public:
   using Storage =
-      llvm::PointerIntPair<llvm::PointerUnion<ValueDecl*, OpaqueValueExpr*>, 2,
+      llvm::PointerIntPair<llvm::PointerUnion<ValueDecl *, Expr *>, 2,
                            unsigned>;
 
 private:
@@ -78,15 +81,7 @@ class CapturedValue {
   CapturedValue(ValueDecl *Val, unsigned Flags, SourceLoc Loc)
       : Value(Val, Flags), Loc(Loc) {}
 
-private:
-  // This is only used in TypeLowering when forming Lowered Capture
-  // Info. OpaqueValueExpr captured value should never show up in the AST
-  // itself.
-  //
-  // NOTE: AbstractClosureExpr::getIsolationCrossing relies upon this and
-  // asserts that it never sees one of these.
-  explicit CapturedValue(OpaqueValueExpr *Val, unsigned Flags)
-      : Value(Val, Flags), Loc(SourceLoc()) {}
+  CapturedValue(Expr *Val, unsigned Flags);
 
 public:
   static CapturedValue getDynamicSelfMetadata() {
@@ -97,10 +92,14 @@ class CapturedValue {
   bool isNoEscape() const { return Value.getInt() & IsNoEscape; }
 
   bool isDynamicSelfMetadata() const { return !Value.getPointer(); }
-  bool isOpaqueValue() const {
-    return Value.getPointer().is<OpaqueValueExpr *>();
+
+  bool isExpr() const {
+    return Value.getPointer().dyn_cast<Expr *>();
   }
 
+  bool isPackElement() const;
+  bool isOpaqueValue() const;
+
   /// Returns true if this captured value is a local capture.
   ///
   /// NOTE: This implies that the value is not dynamic self metadata, since
@@ -116,17 +115,19 @@ class CapturedValue {
   }
 
   ValueDecl *getDecl() const {
-    assert(Value.getPointer() && "dynamic Self metadata capture does not "
-           "have a value");
     return Value.getPointer().dyn_cast<ValueDecl *>();
   }
 
-  OpaqueValueExpr *getOpaqueValue() const {
-    assert(Value.getPointer() && "dynamic Self metadata capture does not "
-           "have a value");
-    return Value.getPointer().dyn_cast<OpaqueValueExpr *>();
+  Expr *getExpr() const {
+    return Value.getPointer().dyn_cast<Expr *>();
   }
 
+  OpaqueValueExpr *getOpaqueValue() const;
+
+  PackElementExpr *getPackElement() const;
+
+  Type getPackElementType() const;
+
   SourceLoc getLoc() const { return Loc; }
 
   unsigned getFlags() const { return Value.getInt(); }

lib/AST/CaptureInfo.cpp

@@ -13,11 +13,36 @@
 #include "swift/AST/CaptureInfo.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
 
+CapturedValue::CapturedValue(Expr *Val, unsigned Flags)
+    : Value(Val, Flags), Loc(SourceLoc()) {
+  assert(isa<OpaqueValueExpr>(Val) || isa<PackElementExpr>(Val));
+}
+
+bool CapturedValue::isPackElement() const {
+  return isExpr() && isa<PackElementExpr>(getExpr());
+}
+bool CapturedValue::isOpaqueValue() const {
+  return isExpr() && isa<OpaqueValueExpr>(getExpr());
+}
+
+OpaqueValueExpr *CapturedValue::getOpaqueValue() const {
+  return dyn_cast_or_null<OpaqueValueExpr>(getExpr());
+}
+
+PackElementExpr *CapturedValue::getPackElement() const {
+  return dyn_cast_or_null<PackElementExpr>(getExpr());
+}
+
+Type CapturedValue::getPackElementType() const {
+  return getPackElement()->getType();
+}
+
 ArrayRef<CapturedValue>
 CaptureInfo::CaptureInfoStorage::getCaptures() const {
   return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), NumCapturedValues);
@@ -126,7 +151,18 @@ void CaptureInfo::print(raw_ostream &OS) const {
 
   interleave(getCaptures(),
              [&](const CapturedValue &capture) {
-               OS << capture.getDecl()->getBaseName();
+               if (capture.getDecl())
+                 OS << capture.getDecl()->getBaseName();
+               else if (capture.isPackElement()) {
+                 OS << "[pack element] ";
+                 capture.getPackElement()->dump(OS);
+               } else if (capture.isOpaqueValue()) {
+                 OS << "[opaque] ";
+                 capture.getOpaqueValue()->dump(OS);
+               } else {
+                 OS << "[unknown] ";
+                 assert(false);
+               }
 
                if (capture.isDirect())
                  OS << "<direct>";