Update type checking for mutate accessors

Author: meg-gupta

include/swift/Sema/SyntacticElementTarget.h

@@ -262,8 +262,9 @@ class SyntacticElementTarget {
                     unsigned patternBindingIndex, bool bindPatternVarsOneWay);
 
   /// Form an expression target for a ReturnStmt.
-  static SyntacticElementTarget
-  forReturn(ReturnStmt *returnStmt, Type contextTy, DeclContext *dc);
+  static SyntacticElementTarget forReturn(ReturnStmt *returnStmt,
+                                          Expr *returnExpr, Type contextTy,
+                                          DeclContext *dc);
 
   /// Form a target for the preamble of a for-in loop, excluding its where
   /// clause and body.

lib/AST/ASTVerifier.cpp

@@ -1097,9 +1097,15 @@ class Verifier : public ASTWalker {
     void verifyChecked(ReturnStmt *S) {
       auto func = Functions.back();
       Type resultType;
+      bool hasInOutResult = false;
+
       if (auto *FD = dyn_cast<FuncDecl>(func)) {
         resultType = FD->getResultInterfaceType();
         resultType = FD->mapTypeIntoContext(resultType);
+        hasInOutResult = FD->getInterfaceType()
+                             ->castTo<AnyFunctionType>()
+                             ->getExtInfo()
+                             .hasInOutResult();
       } else if (auto closure = dyn_cast<AbstractClosureExpr>(func)) {
         resultType = closure->getResultType();
       } else if (isa<ConstructorDecl>(func)) {
@@ -1112,6 +1118,9 @@ class Verifier : public ASTWalker {
         auto result = S->getResult();
         auto returnType = result->getType();
         // Make sure that the return has the same type as the function.
+        if (hasInOutResult) {
+          resultType = InOutType::get(resultType);
+        }
         checkSameType(resultType, returnType, "return type");
       } else {
         // Make sure that the function has a Void result type.

lib/Sema/SyntacticElementTarget.cpp

@@ -170,13 +170,14 @@ SyntacticElementTarget SyntacticElementTarget::forInitialization(
   return result;
 }
 
-SyntacticElementTarget
-SyntacticElementTarget::forReturn(ReturnStmt *returnStmt, Type contextTy,
-                                  DeclContext *dc) {
+SyntacticElementTarget SyntacticElementTarget::forReturn(ReturnStmt *returnStmt,
+                                                         Expr *returnExpr,
+                                                         Type contextTy,
+                                                         DeclContext *dc) {
   assert(contextTy);
   assert(returnStmt->hasResult() && "Must have result to be type-checked");
   ContextualTypeInfo contextInfo(contextTy, CTP_ReturnStmt);
-  SyntacticElementTarget target(returnStmt->getResult(), dc, contextInfo,
+  SyntacticElementTarget target(returnExpr, dc, contextInfo,
                                 /*isDiscarded*/ false);
   target.expression.parentReturnStmt = returnStmt;
   return target;

lib/Sema/TypeCheckDecl.cpp

@@ -2571,11 +2571,6 @@ InterfaceTypeRequest::evaluate(Evaluator &eval, ValueDecl *D) const {
             infoBuilder.withLifetimeDependencies(*lifetimeDependenceInfo);
       }
 
-      auto *accessor = dyn_cast<AccessorDecl>(AFD);
-      if (accessor && accessor->isMutateAccessor()) {
-        infoBuilder = infoBuilder.withHasInOutResult();
-      }
-
       auto info = infoBuilder.build();
 
       if (sig && !hasSelf) {
@@ -2595,7 +2590,10 @@ InterfaceTypeRequest::evaluate(Evaluator &eval, ValueDecl *D) const {
         selfInfoBuilder =
             selfInfoBuilder.withLifetimeDependencies(*lifetimeDependenceInfo);
       }
-
+      auto *accessor = dyn_cast<AccessorDecl>(AFD);
+      if (accessor && accessor->isMutateAccessor()) {
+        selfInfoBuilder = selfInfoBuilder.withHasInOutResult();
+      }
       // FIXME: Verify ExtInfo state is correct, not working by accident.
       auto selfInfo = selfInfoBuilder.build();
       if (sig) {

lib/Sema/TypeCheckStmt.cpp

@@ -1112,8 +1112,26 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       return RS;
     }
 
+    auto *accessor = TheFunc->getAccessorDecl();
+    auto *exprToCheck = RS->getResult();
+    InOutExpr *inout = nullptr;
+
+    if (accessor && accessor->isMutateAccessor()) {
+      // Check that the returned expression is a &.
+      if ((inout = dyn_cast<InOutExpr>(exprToCheck))) {
+        ResultTy = InOutType::get(ResultTy);
+      } else {
+        getASTContext()
+            .Diags
+            .diagnose(exprToCheck->getLoc(), diag::missing_address_of_return)
+            .highlight(exprToCheck->getSourceRange());
+        inout = new (getASTContext()) InOutExpr(
+            exprToCheck->getStartLoc(), exprToCheck, Type(), /*implicit*/ true);
+      }
+    }
     using namespace constraints;
-    auto target = SyntacticElementTarget::forReturn(RS, ResultTy, DC);
+    auto target =
+        SyntacticElementTarget::forReturn(RS, exprToCheck, ResultTy, DC);
     auto resultTarget = TypeChecker::typeCheckTarget(target);
     if (resultTarget) {
       RS->setResult(resultTarget->getAsExpr());