[WIP | SE] [Parse][Sema] Optional iteration

include/swift/AST/DiagnosticsSema.def

@@ -3142,6 +3142,13 @@ ERROR(continue_outside_loop,none,
 ERROR(continue_not_in_this_stmt,none,
       "'continue' cannot be used with %0 statements", (StringRef))
 
+ERROR(foreach_sequence_not_optional,none,
+      "optional for-in loop must not be used on non-optional sequence of type %0",
+      (Type))
+NOTE(to_optional_foreach_fixit,none,
+     "use 'for?' to ignore the loop when the optional value contains 'nil'", ())
+NOTE(to_nonoptional_foreach_fixit,none, "", ())
+
 // Switch Stmt
 ERROR(no_match_operator,none,
       "no binary '~=' operator available for 'switch' statement", ())

include/swift/AST/Stmt.h

@@ -798,6 +798,7 @@ class ForEachStmt : public LabeledStmt {
   Expr *Sequence;
   Expr *WhereExpr = nullptr;
   BraceStmt *Body;
+  bool isOptional;
 
   /// The iterator variable along with its initializer.
   PatternBindingDecl *Iterator = nullptr;
@@ -808,11 +809,11 @@ class ForEachStmt : public LabeledStmt {
 public:
   ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, Pattern *Pat,
               SourceLoc InLoc, Expr *Sequence, Expr *WhereExpr, BraceStmt *Body,
-              Optional<bool> implicit = None)
+              bool isOptional, Optional<bool> implicit = None)
     : LabeledStmt(StmtKind::ForEach, getDefaultImplicitFlag(implicit, ForLoc),
                   LabelInfo),
       ForLoc(ForLoc), Pat(nullptr), InLoc(InLoc), Sequence(Sequence),
-      WhereExpr(WhereExpr), Body(Body) {
+      WhereExpr(WhereExpr), Body(Body), isOptional(isOptional) {
     setPattern(Pat);
   }
 
@@ -821,7 +822,15 @@ class ForEachStmt : public LabeledStmt {
 
   /// getInLoc - Retrieve the location of the 'in' keyword.
   SourceLoc getInLoc() const { return InLoc; }
-
+
+  /// Retrieve the location of '?' in 'for?', present when the iteration
+  /// is optional.
+  SourceLoc getQuestionLoc() const {
+    return ForLoc.isInvalid() ? ForLoc : ForLoc.getAdvancedLoc(3);
+  }
+
+  bool getIsOptional() const { return isOptional; }
+
   /// getPattern - Retrieve the pattern describing the iteration variables.
   /// These variables will only be visible within the body of the loop.
   Pattern *getPattern() const { return Pat; }

lib/Parse/ParseStmt.cpp

@@ -2003,21 +2003,25 @@ static bool isStmtForCStyle(Parser &P) {
 
 /// 
 ///   stmt-for-each:
-///     (identifier ':')? 'for' pattern 'in' expr-basic \
+///     (identifier ':')? 'for'(?)? pattern 'in' expr-basic \
 ///             ('where' expr-basic)? stmt-brace
 ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
   SyntaxContext->setCreateSyntax(SyntaxKind::ForInStmt);
   SourceLoc ForLoc = consumeToken(tok::kw_for);
   ParserStatus Status;
   ParserResult<Pattern> pattern;
   ParserResult<Expr> Container;
+  bool isOptional = false;
 
   // The C-style for loop which was supported in Swift2 and foreach-style-for
   // loop are conflated together into a single keyword, so we have to do some
   // lookahead to resolve what is going on.
   bool IsCStyleFor = isStmtForCStyle(*this);
   auto StartOfControl = Tok.getLoc();
 
+  if (consumeIf(tok::question_postfix))
+    isOptional = true;
+
   // Parse the pattern.  This is either 'case <refutable pattern>' or just a
   // normal pattern.
   if (consumeIf(tok::kw_case)) {
@@ -2122,9 +2126,9 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
 
   return makeParserResult(
       Status,
-      new (Context) ForEachStmt(LabelInfo, ForLoc, pattern.get(), InLoc,
-                                Container.get(), Where.getPtrOrNull(),
-                                Body.get()));
+      new (Context) ForEachStmt(LabelInfo, ForLoc, pattern.get(),
+                                InLoc, Container.get(), Where.getPtrOrNull(),
+                                Body.get(), isOptional));
 }
 
 ///

lib/Sema/CSApply.cpp

@@ -7868,7 +7868,7 @@ Expr *TypeChecker::callWitness(Expr *base, DeclContext *dc,
 
   if (auto metaType = type->getAs<AnyMetatypeType>())
     type = metaType->getInstanceType();
-  
+
   auto witness = findNamedWitnessImpl(
                    *this, dc, type->getRValueType(), protocol,
                    name, brokenProtocolDiag);
@@ -7972,6 +7972,26 @@ Expr *TypeChecker::callWitness(Expr *base, DeclContext *dc,
   if (!result)
     return nullptr;
 
+  if (dyn_cast<BindOptionalExpr>(base)) {
+    Expr *OEE = nullptr;
+    if (cs.getType(call)->getOptionalObjectType()) {
+      OEE = new (Context) OptionalEvaluationExpr(call);
+      OEE->setType(cs.getType(call));
+    } else {
+      Expr *IIO = new (Context) InjectIntoOptionalExpr(
+          call, OptionalType::get(cs.getType(call)));
+      OEE = new (Context) OptionalEvaluationExpr(IIO);
+      OEE->setType(IIO->getType());
+    }
+    OEE->print(llvm::outs());llvm::outs()<<"\n\n&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&\n\n";
+//    auto result = rewriter.visitOptionalEvaluationExpr(oee);
+
+//    if (!OEE)
+//      return nullptr;
+//    rewriter.finalize(oee);
+    return OEE;
+  }
+
   rewriter.finalize(result);
   return result;
 }

lib/Sema/CSDiag.cpp

@@ -8110,6 +8110,7 @@ bool FailureDiagnosis::diagnoseExprFailure() {
   // be type checked on its own (even to an incomplete type) then that is where
   // we focus our attention.  If we do find a type, we use it to check for
   // contextual type mismatches.
+  expr->print(llvm::outs());llvm::outs()<<" ±±±±±±±±±±±±±±±±±±±±±±±±±±±±±±±\n";
   return visit(expr);
 }
 

lib/Sema/CSDiagnostics.cpp

@@ -407,62 +407,6 @@ bool MemberAccessOnOptionalBaseFailure::diagnoseAsError() {
                                            resultIsOptional, SourceRange());
 }
 
-// Suggest a default value via ?? <default value>
-static void offerDefaultValueUnwrapFixit(TypeChecker &TC, DeclContext *DC, Expr *expr) {
-  auto diag =
-  TC.diagnose(expr->getLoc(), diag::unwrap_with_default_value);
-
-  // Figure out what we need to parenthesize.
-  bool needsParensInside =
-  exprNeedsParensBeforeAddingNilCoalescing(TC, DC, expr);
-  bool needsParensOutside =
-  exprNeedsParensAfterAddingNilCoalescing(TC, DC, expr, expr);
-
-  llvm::SmallString<2> insertBefore;
-  llvm::SmallString<32> insertAfter;
-  if (needsParensOutside) {
-    insertBefore += "(";
-  }
-  if (needsParensInside) {
-    insertBefore += "(";
-    insertAfter += ")";
-  }
-  insertAfter += " ?? <" "#default value#" ">";
-  if (needsParensOutside)
-    insertAfter += ")";
-
-  if (!insertBefore.empty()) {
-    diag.fixItInsert(expr->getStartLoc(), insertBefore);
-  }
-  diag.fixItInsertAfter(expr->getEndLoc(), insertAfter);
-}
-
-// Suggest a force-unwrap.
-static void offerForceUnwrapFixit(ConstraintSystem &CS, Expr *expr) {
-  auto diag = CS.TC.diagnose(expr->getLoc(), diag::unwrap_with_force_value);
-
-  // If expr is optional as the result of an optional chain and this last
-  // dot isn't a member returning optional, then offer to force the last
-  // link in the chain, rather than an ugly parenthesized postfix force.
-  if (auto optionalChain = dyn_cast<OptionalEvaluationExpr>(expr)) {
-    if (auto dotExpr =
-        dyn_cast<UnresolvedDotExpr>(optionalChain->getSubExpr())) {
-      auto bind = dyn_cast<BindOptionalExpr>(dotExpr->getBase());
-      if (bind && !CS.getType(dotExpr)->getOptionalObjectType()) {
-        diag.fixItReplace(SourceRange(bind->getLoc()), "!");
-        return;
-      }
-    }
-  }
-
-  if (expr->canAppendPostfixExpression(true)) {
-    diag.fixItInsertAfter(expr->getEndLoc(), "!");
-  } else {
-    diag.fixItInsert(expr->getStartLoc(), "(")
-    .fixItInsertAfter(expr->getEndLoc(), ")!");
-  }
-}
-
 class VarDeclMultipleReferencesChecker : public ASTWalker {
   VarDecl *varDecl;
   int count;
@@ -488,6 +432,10 @@ static bool diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type type) {
   CS.TC.diagnose(expr->getLoc(), diag::optional_not_unwrapped, type,
                  unwrappedType);
 
+  auto getType = [&](const Expr *E) -> Type {
+    return CS.getType(E);
+  };
+
   // If the expression we're unwrapping is the only reference to a
   // local variable whose type isn't explicit in the source, then
   // offer unwrapping fixits on the initializer as well.
@@ -530,13 +478,13 @@ static bool diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type type) {
 
         offerDefaultValueUnwrapFixit(CS.TC, varDecl->getDeclContext(),
                                      initializer);
-        offerForceUnwrapFixit(CS, initializer);
+        offerForceUnwrapFixit(initializer, CS.TC, getType);
       }
     }
   }
 
   offerDefaultValueUnwrapFixit(CS.TC, CS.DC, expr);
-  offerForceUnwrapFixit(CS, expr);
+  offerForceUnwrapFixit(expr, CS.TC, getType);
   return true;
 }
 

lib/Sema/MiscDiagnostics.cpp

@@ -16,6 +16,7 @@
 
 #include "MiscDiagnostics.h"
 #include "TypeChecker.h"
+#include "ConstraintSystem.h"
 #include "TypeCheckAvailability.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/NameLookup.h"
@@ -2979,6 +2980,33 @@ static void checkSwitch(TypeChecker &TC, const SwitchStmt *stmt) {
   }
 }
 
+static void checkForEach(TypeChecker &TC, DeclContext *DC,
+                         const ForEachStmt *stmt) {
+  const auto seq = stmt->getSequence();
+  const auto seqTy = seq->getType();
+
+  auto unwrappedTy = seqTy->getOptionalObjectType();
+
+  if (unwrappedTy) {
+    if (!stmt->getIsOptional()) {
+      TC.diagnose(seq->getLoc(), diag::optional_not_unwrapped, seqTy,
+                  unwrappedTy);
+
+      TC.diagnose(seq->getLoc(), diag::to_optional_foreach_fixit)
+        .fixItInsert(stmt->getQuestionLoc(),
+                     getTokenText(tok::question_postfix));
+
+      offerDefaultValueUnwrapFixit(TC, DC, seq);
+      offerForceUnwrapFixit(seq, TC);
+    }
+  } else if (stmt->getIsOptional()) {
+    TC.diagnose(stmt->getForLoc(), diag::foreach_sequence_not_optional,
+                seqTy)
+      .highlight(seq->getSourceRange())
+      .fixItRemove(stmt->getQuestionLoc());
+  }
+}
+
 void swift::fixItEncloseTrailingClosure(TypeChecker &TC,
                                         InFlightDiagnostic &diag,
                                         const CallExpr *call,
@@ -3882,6 +3910,62 @@ static void diagnoseDeprecatedWritableKeyPath(TypeChecker &TC, const Expr *E,
   const_cast<Expr *>(E)->walk(Walker);
 }
 
+void swift::offerDefaultValueUnwrapFixit(TypeChecker &TC, DeclContext *DC,
+                                         Expr *expr) {
+  auto diag =
+  TC.diagnose(expr->getLoc(), diag::unwrap_with_default_value);
+
+  // Figure out what we need to parenthesize.
+  bool needsParensInside =
+    exprNeedsParensBeforeAddingNilCoalescing(TC, DC, expr);
+  bool needsParensOutside =
+    exprNeedsParensAfterAddingNilCoalescing(TC, DC, expr, expr);
+
+  llvm::SmallString<2> insertBefore;
+  llvm::SmallString<32> insertAfter;
+  if (needsParensOutside) {
+    insertBefore += "(";
+  }
+  if (needsParensInside) {
+    insertBefore += "(";
+    insertAfter += ")";
+  }
+  insertAfter += " ?? <" "#default value#" ">";
+  if (needsParensOutside)
+    insertAfter += ")";
+
+  if (!insertBefore.empty()) {
+    diag.fixItInsert(expr->getStartLoc(), insertBefore);
+  }
+  diag.fixItInsertAfter(expr->getEndLoc(), insertAfter);
+}
+
+void swift::offerForceUnwrapFixit(Expr *expr, TypeChecker &TC,
+                           llvm::function_ref<Type(const Expr *)> getType) {
+  auto diag = TC.diagnose(expr->getLoc(), diag::unwrap_with_force_value);
+
+  // If expr is optional as the result of an optional chain and this last
+  // dot isn't a member returning optional, then offer to force the last
+  // link in the chain, rather than an ugly parenthesized postfix force.
+  if (auto optionalChain = dyn_cast<OptionalEvaluationExpr>(expr)) {
+    if (auto dotExpr =
+        dyn_cast<UnresolvedDotExpr>(optionalChain->getSubExpr())) {
+      auto bind = dyn_cast<BindOptionalExpr>(dotExpr->getBase());
+      if (bind && !getType(dotExpr)->getOptionalObjectType()) {
+        diag.fixItReplace(SourceRange(bind->getLoc()), "!");
+        return;
+      }
+    }
+  }
+
+  if (expr->canAppendPostfixExpression(true)) {
+    diag.fixItInsertAfter(expr->getEndLoc(), "!");
+  } else {
+    diag.fixItInsert(expr->getStartLoc(), "(")
+    .fixItInsertAfter(expr->getEndLoc(), ")!");
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // High-level entry points.
 //===----------------------------------------------------------------------===//
@@ -3903,11 +3987,14 @@ void swift::performSyntacticExprDiagnostics(TypeChecker &TC, const Expr *E,
     diagDeprecatedObjCSelectors(TC, DC, E);
 }
 
-void swift::performStmtDiagnostics(TypeChecker &TC, const Stmt *S) {
+void swift::performStmtDiagnostics(TypeChecker &TC, DeclContext *DC, const Stmt *S) {
   TC.checkUnsupportedProtocolType(const_cast<Stmt *>(S));
 
-  if (auto switchStmt = dyn_cast<SwitchStmt>(S))
+  if (auto switchStmt = dyn_cast<SwitchStmt>(S)) {
     checkSwitch(TC, switchStmt);
+  } else if (auto forEachStmt = dyn_cast<ForEachStmt>(S)) {
+    checkForEach(TC, DC, forEachStmt);
+  }
 
   checkStmtConditionTrailingClosure(TC, S);
 

lib/Sema/MiscDiagnostics.h

@@ -40,7 +40,7 @@ void performSyntacticExprDiagnostics(TypeChecker &TC, const Expr *E,
                                      bool isExprStmt);
 
 /// \brief Emit diagnostics for a given statement.
-void performStmtDiagnostics(TypeChecker &TC, const Stmt *S);
+void performStmtDiagnostics(TypeChecker &TC, DeclContext *DC, const Stmt *S);
 
 void performAbstractFuncDeclDiagnostics(TypeChecker &TC,
                                         AbstractFunctionDecl *AFD);
@@ -97,6 +97,15 @@ void fixItEncloseTrailingClosure(TypeChecker &TC,
                                  const CallExpr *call,
                                  Identifier closureLabel);
 
+// Suggest a default value via ?? <default value>
+void offerDefaultValueUnwrapFixit(TypeChecker &TC, DeclContext *DC, Expr *expr);
+
+// Suggest a force-unwrap.
+void offerForceUnwrapFixit(Expr *expr, TypeChecker &TC,
+                           llvm::function_ref<Type(const Expr *)> getType =
+                             [](const Expr *E) -> Type {
+                               return E->getType();
+                             });
 } // namespace swift
 
 #endif // SWIFT_SEMA_MISC_DIAGNOSTICS_H