Introduce if/switch expressions

Introduce SingleValueStmtExpr, which allows the
embedding of a statement in an expression context.
This then allows us to parse and type-check `if`
and `switch` statements as expressions, gated
behind the `IfSwitchExpression` experimental
feature for now. In the future,
SingleValueStmtExpr could also be used for e.g
`do` expressions.

For now, only single expression branches are
supported for producing a value from an
`if`/`switch` expression, and each branch is
type-checked independently. A multi-statement
branch may only appear if it ends with a `throw`,
and it may not `break`, `continue`, or `return`.

The placement of `if`/`switch` expressions is also
currently limited by a syntactic use diagnostic.
Currently they're only allowed in bindings,
assignments, throws, and returns. But this could
be lifted in the future if desired.

Author: hamishknight

include/swift/AST/ASTScope.h

@@ -955,6 +955,7 @@ class PatternEntryInitializerScope final : public AbstractPatternEntryScope {
 
 protected:
   bool lookupLocalsOrMembers(DeclConsumer) const override;
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 /// The scope introduced by a conditional clause initializer in an

include/swift/AST/ASTTypeIDZone.def

@@ -26,6 +26,7 @@ SWIFT_TYPEID(Fingerprint)
 SWIFT_TYPEID(GenericSignature)
 SWIFT_TYPEID(ImplicitImportList)
 SWIFT_TYPEID(ImplicitMemberAction)
+SWIFT_TYPEID(IsSingleValueStmtResult)
 SWIFT_TYPEID(ParamSpecifier)
 SWIFT_TYPEID(PropertyWrapperAuxiliaryVariables)
 SWIFT_TYPEID(PropertyWrapperInitializerInfo)

include/swift/AST/ASTTypeIDs.h

@@ -41,6 +41,7 @@ class GenericParamList;
 class GenericSignature;
 class GenericTypeParamType;
 class InfixOperatorDecl;
+class IsSingleValueStmtResult;
 class IterableDeclContext;
 class ModuleDecl;
 struct ImplicitImportList;

include/swift/AST/CASTBridging.h

@@ -169,6 +169,9 @@ void *SwiftVarDecl_create(void *ctx, BridgedIdentifier _Nullable name,
                           void *initExpr, void *loc, _Bool isStatic,
                           _Bool isLet, void *dc);
 
+void *SingleValueStmtExpr_createWithWrappedBranches(void *ctx, void *S,
+                                                    void *DC, _Bool mustBeExpr);
+
 void *IfStmt_create(void *ctx, void *ifLoc, void *cond, void *_Nullable then,
                     void *_Nullable elseLoc, void *_Nullable elseStmt);
 

include/swift/AST/DiagnosticsSema.def

@@ -1080,6 +1080,27 @@ ERROR(ternary_expr_cases_mismatch,none,
       "result values in '? :' expression have mismatching types %0 and %1",
       (Type, Type))
 
+// Statements as expressions
+ERROR(single_value_stmt_branches_mismatch,none,
+      "branches have mismatching types %0 and %1",
+      (Type, Type))
+ERROR(single_value_stmt_out_of_place,none,
+      "'%0' may only be used as expression in return, throw, or as the source "
+      "of an assignment",
+      (StmtKind))
+ERROR(single_value_stmt_must_be_unlabeled,none,
+      "'%0' cannot have a jump label when used as expression",
+      (StmtKind))
+ERROR(if_expr_must_be_syntactically_exhaustive,none,
+      "'if' must have an unconditional 'else' to be used as expression",
+      ())
+ERROR(single_value_stmt_branch_must_end_in_throw,none,
+      "non-expression branch of '%0' expression may only end with a 'throw'",
+      (StmtKind))
+ERROR(cannot_jump_in_single_value_stmt,none,
+      "cannot '%0' in '%1' when used as expression",
+      (StmtKind, StmtKind))
+
 ERROR(did_not_call_function_value,none,
       "function value was used as a property; add () to call it",
       ())

include/swift/AST/Expr.h

@@ -5969,6 +5969,63 @@ class KeyPathDotExpr : public Expr {
   }
 };
 
+/// An expression that may wrap a statement which produces a single value.
+class SingleValueStmtExpr : public Expr {
+public:
+  enum class Kind {
+    If, Switch
+  };
+
+private:
+  Stmt *S;
+  DeclContext *DC;
+
+  SingleValueStmtExpr(Stmt *S, DeclContext *DC)
+      : Expr(ExprKind::SingleValueStmt, /*isImplicit*/ true), S(S), DC(DC) {}
+
+public:
+  /// Creates a new SingleValueStmtExpr wrapping a statement.
+  static SingleValueStmtExpr *create(ASTContext &ctx, Stmt *S, DeclContext *DC);
+
+  /// Creates a new SingleValueStmtExpr wrapping a statement, and recursively
+  /// attempts to wrap any branches of that statement that can become single
+  /// value statement expressions.
+  ///
+  /// If \p mustBeExpr is true, branches will be eagerly wrapped even if they
+  /// may not be valid SingleValueStmtExprs (which Sema will later diagnose).
+  static SingleValueStmtExpr *createWithWrappedBranches(ASTContext &ctx,
+                                                        Stmt *S,
+                                                        DeclContext *DC,
+                                                        bool mustBeExpr);
+
+  /// Attempt to look through valid parent expressions to a child
+  /// SingleValueStmtExpr.
+  static SingleValueStmtExpr *tryDigOutSingleValueStmtExpr(Expr *E);
+
+  /// Retrieve the wrapped statement.
+  Stmt *getStmt() const { return S; }
+  void setStmt(Stmt *newS) { S = newS; }
+
+  /// Retrieve the kind of statement being wrapped.
+  Kind getStmtKind() const;
+
+  /// Retrieve the complete set of branches for the underlying statement.
+  ArrayRef<Stmt *> getBranches(SmallVectorImpl<Stmt *> &scratch) const;
+
+  /// Retrieve the single expression branches of the statement, excluding
+  /// branches that either have multiple expressions, or have statements.
+  ArrayRef<Expr *>
+  getSingleExprBranches(SmallVectorImpl<Expr *> &scratch) const;
+
+  DeclContext *getDeclContext() const { return DC; }
+
+  SourceRange getSourceRange() const;
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::SingleValueStmt;
+  }
+};
+
 /// Expression node that effects a "one-way" constraint in
 /// the constraint system, allowing type information to flow from the
 /// subexpression outward but not the other way.
@@ -6002,6 +6059,10 @@ class TypeJoinExpr final : public Expr,
 
   DeclRefExpr *Var;
 
+  /// If this is joining the expression branches for a SingleValueStmtExpr,
+  /// this holds the expr node. Otherwise, it is \c nullptr.
+  SingleValueStmtExpr *SVE;
+
   size_t numTrailingObjects() const {
     return getNumElements();
   }
@@ -6011,13 +6072,14 @@ class TypeJoinExpr final : public Expr,
   }
 
   TypeJoinExpr(llvm::PointerUnion<DeclRefExpr *, TypeBase *> result,
-               ArrayRef<Expr *> elements);
+               ArrayRef<Expr *> elements, SingleValueStmtExpr *SVE);
 
   static TypeJoinExpr *
   createImpl(ASTContext &ctx,
              llvm::PointerUnion<DeclRefExpr *, TypeBase *> varOrType,
              ArrayRef<Expr *> elements,
-             AllocationArena arena = AllocationArena::Permanent);
+             AllocationArena arena = AllocationArena::Permanent,
+             SingleValueStmtExpr *SVE = nullptr);
 
 public:
   static TypeJoinExpr *
@@ -6032,6 +6094,12 @@ class TypeJoinExpr final : public Expr,
     return createImpl(ctx, joinType.getPointer(), exprs, arena);
   }
 
+  /// Create a join for the branch types of a SingleValueStmtExpr.
+  static TypeJoinExpr *
+  forBranchesOfSingleValueStmtExpr(ASTContext &ctx, Type joinType,
+                                   SingleValueStmtExpr *SVE,
+                                   AllocationArena arena);
+
   SourceLoc getLoc() const { return SourceLoc(); }
   SourceRange getSourceRange() const { return SourceRange(); }
 
@@ -6054,6 +6122,10 @@ class TypeJoinExpr final : public Expr,
     getMutableElements()[i] = E;
   }
 
+  /// If this is joining the expression branches for a SingleValueStmtExpr,
+  /// this returns the expr node. Otherwise, returns \c nullptr.
+  SingleValueStmtExpr *getSingleValueStmtExpr() const { return SVE; }
+
   unsigned getNumElements() const { return Bits.TypeJoinExpr.NumElements; }
 
   static bool classof(const Expr *E) {

include/swift/AST/ExprNodes.def

@@ -204,6 +204,7 @@ EXPR(LazyInitializer, Expr)
 EXPR(EditorPlaceholder, Expr)
 EXPR(ObjCSelector, Expr)
 EXPR(KeyPath, Expr)
+EXPR(SingleValueStmt, Expr)
 UNCHECKED_EXPR(KeyPathDot, Expr)
 UNCHECKED_EXPR(OneWay, Expr)
 EXPR(Tap, Expr)

include/swift/AST/Stmt.h

@@ -35,13 +35,15 @@ class ASTContext;
 class ASTWalker;
 class Decl;
 class DeclContext;
+class Evaluator;
 class Expr;
 class FuncDecl;
 class Pattern;
 class PatternBindingDecl;
 class VarDecl;
 class CaseStmt;
 class DoCatchStmt;
+class IsSingleValueStmtResult;
 class SwitchStmt;
 
 enum class StmtKind {
@@ -133,7 +135,12 @@ class alignas(8) Stmt : public ASTAllocated<Stmt> {
 
   SourceRange getSourceRange() const;
   SourceLoc TrailingSemiLoc;
-  
+
+  /// Whether the statement can produce a single value, and as such may be
+  /// treated as an expression.
+  IsSingleValueStmtResult mayProduceSingleValue(Evaluator &eval) const;
+  IsSingleValueStmtResult mayProduceSingleValue(ASTContext &ctx) const;
+
   /// isImplicit - Determines whether this statement was implicitly-generated,
   /// rather than explicitly written in the AST.
   bool isImplicit() const { return Bits.Stmt.Implicit; }
@@ -204,6 +211,10 @@ class BraceStmt final : public Stmt,
 
   ASTNode findAsyncNode();
 
+  /// If this brace is wrapping a single expression, returns it. Otherwise
+  /// returns \c nullptr.
+  Expr *getSingleExpressionElement() const;
+
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Brace; }
 };
 
@@ -711,7 +722,14 @@ class IfStmt : public LabeledConditionalStmt {
 
   Stmt *getElseStmt() const { return Else; }
   void setElseStmt(Stmt *s) { Else = s; }
-  
+
+  /// Retrieve the complete set of branches for this if statement, including
+  /// else if statements.
+  ArrayRef<Stmt *> getBranches(SmallVectorImpl<Stmt *> &scratch) const;
+
+  /// Whether the if statement has an unconditional \c else.
+  bool isSyntacticallyExhaustive() const;
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::If; }
 };
@@ -1283,7 +1301,10 @@ class SwitchStmt final : public LabeledStmt,
   AsCaseStmtRange getCases() const {
     return AsCaseStmtRange(getRawCases(), AsCaseStmtWithSkippingNonCaseStmts());
   }
-  
+
+  /// Retrieve the complete set of branches for this switch statement.
+  ArrayRef<Stmt *> getBranches(SmallVectorImpl<Stmt *> &scratch) const;
+
   static bool classof(const Stmt *S) {
     return S->getKind() == StmtKind::Switch;
   }

include/swift/AST/TypeCheckRequests.h

@@ -3778,6 +3778,128 @@ class PreCheckReturnStmtRequest
   bool isCached() const { return true; }
 };
 
+/// The result of the query for whether a statement can produce a single value.
+class IsSingleValueStmtResult {
+public:
+  enum class Kind {
+    /// The statement may become a SingleValueStmtExpr.
+    Valid,
+
+    /// There are non-single-expression branches that do not end in a throw.
+    UnterminatedBranches,
+
+    /// The statement is an 'if' statement without an unconditional 'else'.
+    NonExhaustiveIf,
+
+    /// There are no single-expression branches.
+    NoExpressionBranches,
+
+    /// There is an unhandled statement branch. This should only be the case
+    /// for invalid AST.
+    UnhandledStmt,
+
+    /// There was a circular reference when evaluating the request. This can be
+    /// ignored, as we will have already diagnosed it.
+    CircularReference,
+
+    /// There is a 'break' or 'continue' within the statement that prevents it
+    /// from being treated as an expression.
+    InvalidJumps,
+
+    /// The statement has a jump label, which is invalid for an expression.
+    HasLabel
+  };
+
+private:
+  Kind TheKind;
+  TinyPtrVector<Stmt *> InvalidJumps;
+  TinyPtrVector<Stmt *> UnterminatedBranches;
+
+  IsSingleValueStmtResult(Kind kind) : TheKind(kind) {
+    assert(kind != Kind::UnterminatedBranches && kind != Kind::InvalidJumps);
+  }
+
+  IsSingleValueStmtResult(Kind kind, TinyPtrVector<Stmt *> stmts)
+      : TheKind(kind) {
+    switch (kind) {
+    case Kind::UnterminatedBranches: {
+      UnterminatedBranches = std::move(stmts);
+      break;
+    }
+    case Kind::InvalidJumps: {
+      InvalidJumps = std::move(stmts);
+      break;
+    }
+    default:
+      llvm_unreachable("Unhandled case in switch!");
+    }
+  }
+
+public:
+  static IsSingleValueStmtResult valid() {
+    return IsSingleValueStmtResult(Kind::Valid);
+  }
+  static IsSingleValueStmtResult
+  unterminatedBranches(TinyPtrVector<Stmt *> branches) {
+    return IsSingleValueStmtResult(Kind::UnterminatedBranches,
+                                   std::move(branches));
+  }
+  static IsSingleValueStmtResult nonExhaustiveIf() {
+    return IsSingleValueStmtResult(Kind::NonExhaustiveIf);
+  }
+  static IsSingleValueStmtResult noExpressionBranches() {
+    return IsSingleValueStmtResult(Kind::NoExpressionBranches);
+  }
+  static IsSingleValueStmtResult unhandledStmt() {
+    return IsSingleValueStmtResult(Kind::UnhandledStmt);
+  }
+  static IsSingleValueStmtResult circularReference() {
+    return IsSingleValueStmtResult(Kind::CircularReference);
+  }
+  static IsSingleValueStmtResult invalidJumps(TinyPtrVector<Stmt *> jumps) {
+    return IsSingleValueStmtResult(Kind::InvalidJumps, std::move(jumps));
+  }
+  static IsSingleValueStmtResult hasLabel() {
+    return IsSingleValueStmtResult(Kind::HasLabel);
+  }
+
+  Kind getKind() const { return TheKind; }
+
+  /// For an unterminated branch kind, retrieves the branch.
+  const TinyPtrVector<Stmt *> &getUnterminatedBranches() const {
+    assert(TheKind == Kind::UnterminatedBranches);
+    return UnterminatedBranches;
+  }
+
+  /// For an invalid jump kind, retrieves the list of invalid jumps.
+  const TinyPtrVector<Stmt *> &getInvalidJumps() const {
+    assert(TheKind == Kind::InvalidJumps);
+    return InvalidJumps;
+  }
+
+  explicit operator bool() const {
+    return TheKind == Kind::Valid;
+  }
+};
+
+/// Computes whether a given statement can be treated as a SingleValueStmtExpr.
+class IsSingleValueStmtRequest
+    : public SimpleRequest<IsSingleValueStmtRequest,
+                           IsSingleValueStmtResult(const Stmt *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  IsSingleValueStmtResult
+  evaluate(Evaluator &evaluator, const Stmt *stmt) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 class GetTypeWrapperInitializer
     : public SimpleRequest<GetTypeWrapperInitializer,
                            ConstructorDecl *(NominalTypeDecl *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -446,6 +446,9 @@ SWIFT_REQUEST(TypeChecker, ContinueTargetRequest,
 SWIFT_REQUEST(TypeChecker, PreCheckReturnStmtRequest,
               Stmt *(ReturnStmt *, DeclContext *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, IsSingleValueStmtRequest,
+              IsSingleValueStmtResult(const Stmt *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapperInitializer,
               ConstructorDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)