[WIP] if/switch expressions

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,29 @@ 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))
+WARNING(redundant_effect_marker_on_single_value_stmt,none,
+        "'%0' on an '%1' expression has no effect", (StringRef, StmtKind))
+
 ERROR(did_not_call_function_value,none,
       "function value was used as a property; add () to call it",
       ())

include/swift/AST/Expr.h

@@ -5971,6 +5971,58 @@ 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);
+  /// 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.

include/swift/AST/ExprNodes.def

@@ -205,6 +205,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

@@ -55,6 +55,7 @@ class PropertyWrapperInitializerInfo;
 struct PropertyWrapperLValueness;
 struct PropertyWrapperMutability;
 class RequirementRepr;
+class ReturnStmt;
 class SpecializeAttr;
 class TrailingWhereClause;
 class TypeAliasDecl;
@@ -3749,6 +3750,146 @@ class ContinueTargetRequest
   bool isCached() const { return true; }
 };
 
+/// Precheck a ReturnStmt, which involves some initial validation, as well as
+/// applying a conversion to a FailStmt if needed.
+class PreCheckReturnStmtRequest
+    : public SimpleRequest<PreCheckReturnStmtRequest,
+                           Stmt *(ReturnStmt *, DeclContext *DC),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  Stmt *evaluate(Evaluator &evaluator, ReturnStmt *RS, DeclContext *DC) const;
+
+public:
+  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

@@ -440,6 +440,12 @@ SWIFT_REQUEST(TypeChecker, BreakTargetRequest,
 SWIFT_REQUEST(TypeChecker, ContinueTargetRequest,
               LabeledStmt *(const ContinueStmt *),
               Cached, NoLocationInfo)
+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)

include/swift/Basic/Features.def

@@ -143,6 +143,9 @@ EXPERIMENTAL_FEATURE(ParserValidation, false)
 /// parser.
 EXPERIMENTAL_FEATURE(ParserDiagnostics, false)
 
+/// Whether to allow 'if' and 'switch' statements to be treated as expressions.
+EXPERIMENTAL_FEATURE(StatementExpressions, false)
+
 /// Enables implicit some while also enabling existential `any`
 EXPERIMENTAL_FEATURE(ImplicitSome, false)
 

include/swift/Parse/Parser.h

@@ -49,6 +49,7 @@ namespace swift {
   class RequirementRepr;
   class SILParserStateBase;
   class ScopeInfo;
+  class SingleValueStmtExpr;
   class SourceManager;
   class TupleType;
   class TypeLoc;