Fix Fix-It locations for "unsafe" insertion

lib/Sema/TypeCheckEffects.cpp

@@ -3167,12 +3167,13 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
   /// anchor expression, so we can emit diagnostics at the end.
   llvm::MapVector<Expr *, std::vector<UnsafeUse>> uncoveredUnsafeUses;
 
-  static bool isEffectAnchor(Expr *e) {
+  static bool isEffectAnchor(Expr *e, bool stopAtAutoClosure) {
     if (e->getLoc().isInvalid())
       return false;
 
-    return isa<AbstractClosureExpr>(e) || isa<DiscardAssignmentExpr>(e) ||
-           isa<AssignExpr>(e) || (isa<DeclRefExpr>(e) && e->isImplicit());
+    return isa<ClosureExpr>(e) || isa<DiscardAssignmentExpr>(e) ||
+           isa<AssignExpr>(e) ||
+           (stopAtAutoClosure && isa<AutoClosureExpr>(e));
   }
 
   static bool isAnchorTooEarly(Expr *e) {
@@ -3181,11 +3182,12 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
 
   /// Find the top location where we should put the await
   static Expr *walkToAnchor(Expr *e, llvm::DenseMap<Expr *, Expr *> &parentMap,
-                            bool isInterpolatedString) {
+                            bool isInterpolatedString,
+                            bool stopAtAutoClosure) {
     llvm::SmallPtrSet<Expr *, 4> visited;
     Expr *parent = e;
     Expr *lastParent = e;
-    while (parent && !isEffectAnchor(parent)) {
+    while (parent && !isEffectAnchor(parent, stopAtAutoClosure)) {
       lastParent = parent;
       parent = parentMap[parent];
 
@@ -3890,7 +3892,8 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
                  Flags.has(ContextFlags::InAsyncLet))) {
         Expr *expr = E.dyn_cast<Expr*>();
         Expr *anchor = walkToAnchor(expr, parentMap,
-                                    CurContext.isWithinInterpolatedString());
+                                    CurContext.isWithinInterpolatedString(),
+                                    /*stopAtAutoClosure=*/true);
         if (Flags.has(ContextFlags::StmtExprCoversAwait))
           classification.setDowngradeToWarning(true);
         if (uncoveredAsync.find(anchor) == uncoveredAsync.end())
@@ -3915,7 +3918,8 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
       if (!Flags.has(ContextFlags::IsUnsafeCovered)) {
         Expr *expr = E.dyn_cast<Expr*>();
         Expr *anchor = walkToAnchor(expr, parentMap,
-                                    CurContext.isWithinInterpolatedString());
+                                    CurContext.isWithinInterpolatedString(),
+                                    /*stopAtAutoClosure=*/false);
 
         // Figure out a location to use if the unsafe use didn't have one.
         SourceLoc replacementLoc;

test/Unsafe/safe.swift

@@ -140,7 +140,7 @@ func testKeyPath() {
 func takesAutoclosure<T>(_ body: @autoclosure () -> T) { }
 
 func testAutoclosure() {
-  // expected-warning@+1{{expression uses unsafe constructs but is not marked with 'unsafe'}}{{20-20=unsafe }}
+  // expected-warning@+1{{expression uses unsafe constructs but is not marked with 'unsafe'}}{{3-3=unsafe }}
   takesAutoclosure(unsafeFunction()) // expected-note{{reference to unsafe global function 'unsafeFunction()'}}
 
   unsafe takesAutoclosure(unsafeFunction())

test/Unsafe/unsafe.swift

@@ -123,6 +123,26 @@ class ExclusivityChecking {
 // TODO: diagnose the need for @unsafe when there's an unsafe superclass.
 class UnsafeSub: UnsafeSuper { }
 
+// -----------------------------------------------------------------------
+// Miscellaneous expression issues
+// -----------------------------------------------------------------------
+
+struct BufferThingy<T> {
+  @unsafe init(count: Int) { }
+}
+
+func testConstruction() {
+  let _ = BufferThingy<Int>(count: 17) // expected-warning{{expression uses unsafe constructs but is not marked with 'unsafe' [Unsafe]}}{{11-11=unsafe }}
+  // expected-note@-1{{reference to unsafe initializer 'init(count:)'}}
+}
+
+func testRHS(b: Bool, x: Int) {
+  @unsafe let limit = 17
+  if b && x > limit { // expected-warning{{expression uses unsafe constructs but is not marked with 'unsafe' [Unsafe]}}{{6-6=unsafe }}
+    // expected-note@-1{{reference to unsafe let 'limit'}}
+  }
+}
+
 // -----------------------------------------------------------------------
 // Declaration references
 // -----------------------------------------------------------------------