[AST] Expand TypeJoin expression to support joining over a type

(cherry picked from commit 4efc35a)

Author: xedin

include/swift/AST/Expr.h

@@ -6006,11 +6006,24 @@ class TypeJoinExpr final : public Expr,
     return { getTrailingObjects<Expr *>(), getNumElements() };
   }
 
-  TypeJoinExpr(DeclRefExpr *var, ArrayRef<Expr *> elements);
+  TypeJoinExpr(llvm::PointerUnion<DeclRefExpr *, TypeBase *> result,
+               ArrayRef<Expr *> elements);
+
+  static TypeJoinExpr *
+  createImpl(ASTContext &ctx,
+             llvm::PointerUnion<DeclRefExpr *, TypeBase *> varOrType,
+             ArrayRef<Expr *> elements);
 
 public:
   static TypeJoinExpr *create(ASTContext &ctx, DeclRefExpr *var,
-                              ArrayRef<Expr *> exprs);
+                              ArrayRef<Expr *> exprs) {
+    return createImpl(ctx, var, exprs);
+  }
+
+  static TypeJoinExpr *create(ASTContext &ctx, Type joinType,
+                              ArrayRef<Expr *> exprs) {
+    return createImpl(ctx, joinType.getPointer(), exprs);
+  }
 
   SourceLoc getLoc() const { return SourceLoc(); }
   SourceRange getSourceRange() const { return SourceRange(); }

lib/AST/ASTDumper.cpp

@@ -2981,8 +2981,12 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
   void visitTypeJoinExpr(TypeJoinExpr *E) {
     printCommon(E, "type_join_expr");
 
-    PrintWithColorRAII(OS, DeclColor) << " var=";
-    printRec(E->getVar());
+    if (auto *var = E->getVar()) {
+      PrintWithColorRAII(OS, DeclColor) << " var=";
+      printRec(var);
+      OS << '\n';
+    }
+
     OS << '\n';
 
     for (auto *member : E->getElements()) {

lib/AST/ASTWalker.cpp

@@ -1260,10 +1260,12 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitTypeJoinExpr(TypeJoinExpr *E) {
-    if (auto *newVar = dyn_cast<DeclRefExpr>(doIt(E->getVar()))) {
-      E->setVar(newVar);
-    } else {
-      return nullptr;
+    if (auto *var = E->getVar()) {
+      if (auto *newVar = dyn_cast<DeclRefExpr>(doIt(var))) {
+        E->setVar(newVar);
+      } else {
+        return nullptr;
+      }
     }
 
     for (unsigned i = 0, e = E->getNumElements(); i != e; ++i) {

lib/AST/Expr.cpp

@@ -2460,21 +2460,31 @@ RegexLiteralExpr::createParsed(ASTContext &ctx, SourceLoc loc,
                                     /*implicit*/ false);
 }
 
-TypeJoinExpr::TypeJoinExpr(DeclRefExpr *varRef, ArrayRef<Expr *> elements)
-  : Expr(ExprKind::TypeJoin, /*implicit=*/true, Type()), Var(varRef) {
-  assert(Var);
+TypeJoinExpr::TypeJoinExpr(llvm::PointerUnion<DeclRefExpr *, TypeBase *> result,
+                           ArrayRef<Expr *> elements)
+    : Expr(ExprKind::TypeJoin, /*implicit=*/true, Type()), Var(nullptr) {
+
+  if (auto *varRef = result.dyn_cast<DeclRefExpr *>()) {
+    assert(varRef);
+    Var = varRef;
+  } else {
+    auto joinType = Type(result.get<TypeBase *>());
+    assert(joinType && "expected non-null type");
+    setType(joinType);
+  }
 
   Bits.TypeJoinExpr.NumElements = elements.size();
   // Copy elements.
   std::uninitialized_copy(elements.begin(), elements.end(),
                           getTrailingObjects<Expr *>());
 }
 
-TypeJoinExpr *TypeJoinExpr::create(ASTContext &ctx, DeclRefExpr *var,
-                                   ArrayRef<Expr *> elements) {
+TypeJoinExpr *TypeJoinExpr::createImpl(
+    ASTContext &ctx, llvm::PointerUnion<DeclRefExpr *, TypeBase *> varOrType,
+    ArrayRef<Expr *> elements) {
   size_t size = totalSizeToAlloc<Expr *>(elements.size());
   void *mem = ctx.Allocate(size, alignof(TypeJoinExpr));
-  return new (mem) TypeJoinExpr(var, elements);
+  return new (mem) TypeJoinExpr(varOrType, elements);
 }
 
 SourceRange MacroExpansionExpr::getSourceRange() const {

lib/Sema/CSGen.cpp

@@ -3692,7 +3692,16 @@ namespace {
                               CS.getConstraintLocator(element));
       }
 
-      auto resultTy = CS.getType(expr->getVar());
+      Type resultTy;
+
+      if (auto *var = expr->getVar()) {
+        resultTy = CS.getType(var);
+      } else {
+        resultTy = expr->getType();
+      }
+
+      assert(resultTy);
+
       // The type of a join expression is obtained by performing
       // a "join-meet" operation on deduced types of its elements
       // and the underlying variable.

lib/Sema/CSSyntacticElement.cpp

@@ -61,6 +61,14 @@ class TypeVariableRefFinder : public ASTWalker {
       ClosureDCs.push_back(closure);
     }
 
+    if (auto *joinExpr = dyn_cast<TypeJoinExpr>(expr)) {
+      // If this join is over a known type, let's
+      // analyze it too because it can contain type
+      // variables.
+      if (!joinExpr->getVar())
+        inferVariables(joinExpr->getType());
+    }
+
     if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
       auto *decl = DRE->getDecl();