Visited single expression of AbstractFunctionDecl.

Taught getSingleExpressionBody to return null and
setSingleExpressionBody to require null in the case that there is a
single stmt of type FailStmt which replaces return nil during sema for
single expression failable initializers (i.e. init?() { nil }).

lib/AST/ASTWalker.cpp

@@ -304,6 +304,15 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     PrettyStackTraceDecl debugStack("walking into body of", AFD);
 #endif
 
+    if (AFD->hasSingleExpressionBody()) {
+      if (auto body = AFD->getSingleExpressionBody()) {
+        if ((body = doIt(body)))
+          AFD->setSingleExpressionBody(body);
+        else
+          return true;
+      }
+    }
+
     bool WalkGenerics = AFD->getGenericParams() &&
         Walker.shouldWalkIntoGenericParams() &&
         // accessor generics are visited from the storage decl
@@ -330,15 +339,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       }
     }
 
-    // FIXME: FIXME_NOW: DETERMINE: Is this or something like it needed?  Right
-    //        now it's causing assertion failures during AST verification.
-    //     
-    //     if (AFD->hasSingleExpressionBody()) {
-    //       if (Expr *body = doIt(AFD->getSingleExpressionBody()))
-    //         AFD->setSingleExpressionBody(body);
-    //       else
-    //         return true;
-    //     }
     if (AFD->getBody(/*canSynthesize=*/false)) {
       AbstractFunctionDecl::BodyKind PreservedKind = AFD->getBodyKind();
       if (BraceStmt *S = cast_or_null<BraceStmt>(doIt(AFD->getBody())))

lib/AST/Decl.cpp

@@ -450,17 +450,37 @@ Expr *AbstractFunctionDecl::getSingleExpressionBody() const {
   auto braceStmt = getBody();
   assert(braceStmt != nullptr && "No body currently available.");
   auto body = getBody()->getElement(0);
-  if (auto *stmt = body.dyn_cast<Stmt *>())
-    return cast<ReturnStmt>(stmt)->getResult();
+  if (auto *stmt = body.dyn_cast<Stmt *>()) {
+    if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
+      return returnStmt->getResult();
+    } else if (auto *failStmt = dyn_cast<FailStmt>(stmt)) {
+      // We can only get to this point if we're a type-checked ConstructorDecl
+      // which was originally spelled init?(...) { nil }.  
+      //
+      // There no longer is a single-expression to return, so ignore null.
+      return nullptr;
+    }
+  }
   return body.get<Expr *>();
 }
 
 void AbstractFunctionDecl::setSingleExpressionBody(Expr *NewBody) {
   assert(hasSingleExpressionBody() && "Not a single-expression body");
   auto body = getBody()->getElement(0);
   if (auto *stmt = body.dyn_cast<Stmt *>()) {
-    cast<ReturnStmt>(stmt)->setResult(NewBody);
-    return;
+    if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
+      returnStmt->setResult(NewBody);
+      return;
+    } else if (auto *failStmt = dyn_cast<FailStmt>(stmt)) {
+      // We can only get to this point if we're a type-checked ConstructorDecl
+      // which was originally spelled init?(...) { nil }.  
+      //
+      // We can no longer write the single-expression which is being set on us 
+      // into anything because a FailStmt does not have such a child.  As a result
+      // we need to demand that the NewBody is null.
+      assert(NewBody == nullptr);
+      return;
+    }
   }
   getBody()->setElement(0, NewBody);
 }