Merge pull request #12671 from hamishknight/subscript-non-escaping-args

[Sema] Allow non-escaping functions to be passed as subscript arguments

Author: slavapestov

lib/Sema/MiscDiagnostics.cpp

@@ -233,6 +233,12 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
         if (auto *IOE = dyn_cast<InOutExpr>(SE->getBase()))
           if (IOE->isImplicit())
             AcceptableInOutExprs.insert(IOE);
+
+        visitIndices(SE, [&](unsigned argIndex, Expr *arg) {
+          arg = lookThroughArgument(arg);
+          if (auto *DRE = dyn_cast<DeclRefExpr>(arg))
+            checkNoEscapeParameterUse(DRE, SE, OperandKind::Argument);
+        });
       }
 
       // Check decl refs in withoutActuallyEscaping blocks.
@@ -352,10 +358,12 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
       return { true, E };
     }
 
-    static void visitArguments(ApplyExpr *apply,
-                               llvm::function_ref<void(unsigned, Expr*)> fn) {
-      auto *arg = apply->getArg();
-
+    /// Visit the argument/s represented by either a ParenExpr or TupleExpr,
+    /// unshuffling if needed. If any other kind of expression, will pass it
+    /// straight back.
+    static void argExprVisitArguments(Expr* arg,
+                                      llvm::function_ref
+                                        <void(unsigned, Expr*)> fn) {
       // The argument could be shuffled if it includes default arguments,
       // label differences, or other exciting things like that.
       if (auto *TSE = dyn_cast<TupleShuffleExpr>(arg))
@@ -373,6 +381,18 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
       }
     }
 
+    static void visitIndices(SubscriptExpr *subscript,
+                             llvm::function_ref<void(unsigned, Expr*)> fn) {
+      auto *indexArgs = subscript->getIndex();
+      argExprVisitArguments(indexArgs, fn);
+    }
+
+    static void visitArguments(ApplyExpr *apply,
+                               llvm::function_ref<void(unsigned, Expr*)> fn) {
+      auto *arg = apply->getArg();
+      argExprVisitArguments(arg, fn);
+    }
+
     static Expr *lookThroughArgument(Expr *arg) {
       while (1) {
         if (auto conv = dyn_cast<ImplicitConversionExpr>(arg))
@@ -655,6 +675,9 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
           if (isa<ParamDecl>(DRE->getDecl()) && useKind == OperandKind::Callee)
             checkNoEscapeParameterCall(apply);
           return;
+        } else if (isa<SubscriptExpr>(parent)
+                   && useKind == OperandKind::Argument) {
+          return;
         } else if (isa<MakeTemporarilyEscapableExpr>(parent)) {
           return;
         }

test/decl/subscript/noescape_accessors.swift

@@ -54,6 +54,15 @@ struct Subscripts {
       global = value
     }
   }
+  
+  subscript(nonescapingIndexWithAddressor fn: () -> Void) -> Int {
+    get {
+      return 0
+    }
+    mutableAddressWithNativeOwner {
+      fatalError()
+    }
+  }
 
   // expected-note@+1 2 {{implicitly non-escaping}}
   subscript(nonescapingIndex fn: () -> ()) -> Int {
@@ -95,14 +104,12 @@ func testSubscripts_value2(nonescaping: () -> (),
   s[value2: 0] = nonescaping // expected-error {{assigning non-escaping parameter}}
 }
 
-// FIXME: Allow these uses in subscript expressions!
-// expected-note@+1 2 {{implicitly non-escaping}}
 func testSubscripts_nonescapingIndex(nonescaping: () -> (),
                                      escaping: @escaping () -> ()) {
   var s = Subscripts()
-  _ = s[nonescapingIndex: nonescaping] // expected-error{{may only be called}}
+  _ = s[nonescapingIndex: nonescaping]
   _ = s[nonescapingIndex: escaping]
-  s[nonescapingIndex: nonescaping] = 0 // expected-error{{may only be called}}
+  s[nonescapingIndex: nonescaping] = 0
   s[nonescapingIndex: escaping] = 0
 }