[DefiniteInitialization] Fix the bug that allows overwriting immutable

variables of optional type (SR-7226, rdar://38624845)

Author: ravikandhadai

lib/SILOptimizer/Mandatory/DIMemoryUseCollectorOwnership.cpp

@@ -950,6 +950,16 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
       continue;
     }
 
+    // unchecked_take_enum_data_addr takes the address of the payload of the
+    // optional, which could be used to update the payload. So, visit the
+    // users of this instruction and ensure that there are no overwrites to an
+    // immutable optional. Note that this special handling is for checking
+    // immutability and is not for checking initialization before use.
+    if (auto *enumDataAddr = dyn_cast<UncheckedTakeEnumDataAddrInst>(User)) {
+      collectUses(enumDataAddr, BaseEltNo);
+      continue;
+    }
+
     // We model destroy_addr as a release of the entire value.
     if (isa<DestroyAddrInst>(User)) {
       trackDestroy(User);

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -1239,18 +1239,14 @@ void LifetimeChecker::handleEscapeUse(const DIMemoryUse &Use) {
     noteUninitializedMembers(Use);
     return;
   }
- 
 
   Diag<StringRef, bool> DiagMessage;
-  if (isa<MarkFunctionEscapeInst>(Inst)) {
-    if (Inst->getLoc().isASTNode<AbstractClosureExpr>())
-      DiagMessage = diag::variable_closure_use_uninit;
-    else
-      DiagMessage = diag::variable_function_use_uninit;
-  } else if (isa<UncheckedTakeEnumDataAddrInst>(Inst)) {
-    DiagMessage = diag::variable_used_before_initialized;
-  } else {
+  if (Inst->getLoc().isASTNode<AbstractClosureExpr>()) {
     DiagMessage = diag::variable_closure_use_uninit;
+  } else if (isa<MarkFunctionEscapeInst>(Inst)) {
+    DiagMessage = diag::variable_function_use_uninit;
+  } else {
+    DiagMessage = diag::variable_used_before_initialized;
   }
 
   diagnoseInitError(Use, DiagMessage);
@@ -1708,10 +1704,11 @@ void LifetimeChecker::handleLoadUseFailure(const DIMemoryUse &Use,
   // If this is a load into a promoted closure capture, diagnose properly as
   // a capture.
   if ((isa<LoadInst>(Inst) || isa<LoadBorrowInst>(Inst)) &&
-      Inst->getLoc().isASTNode<AbstractClosureExpr>())
+      Inst->getLoc().isASTNode<AbstractClosureExpr>()) {
     diagnoseInitError(Use, diag::variable_closure_use_uninit);
-  else
+  } else {
     diagnoseInitError(Use, diag::variable_used_before_initialized);
+  }
 }
 
 /// handleSelfInitUse - When processing a 'self' argument on a class, this is

test/SILOptimizer/definite_init_diagnostics.swift

@@ -1355,3 +1355,130 @@ class ClassWithUnownedProperties {
     c2 = tmp
   }
 }
+
+// Tests for DI when optionals are defined using unchecked_take_enum_data_addr
+// <rdar://38624845>
+
+func testOptionalDoubleWrite() -> String? {
+  let sConst: String? // expected-note {{change 'let' to 'var' to make it mutable}}
+  sConst = ""
+  sConst? = "v2" // expected-error {{immutable value 'sConst' may only be initialized once}}
+  return sConst
+}
+
+func testOptionalDoubleWrite2() -> Int? {
+  let x: Int? // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = 0
+  x? = 0 // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+protocol DIOptionalTestProtocol {
+  var f: Int { get set }
+}
+
+func testOptionalDoubleWrite3(p1: DIOptionalTestProtocol) -> DIOptionalTestProtocol? {
+  let x: DIOptionalTestProtocol? // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = p1
+  x? = p1 // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+func testOptionalWrite() {
+  let x: Int? // expected-note {{constant defined here}}
+              // expected-warning@-1 {{immutable value 'x' was never used; consider removing it}}
+  x? = 0 // expected-error {{constant 'x' used before being initialized}}
+}
+
+func testOptionalWriteGenerics<T>(p: T) -> T? {
+  let x: T? // expected-note {{constant defined here}}
+            // expected-note@-1 {{change 'let' to 'var' to make it mutable}}
+  x? = p  // expected-error {{constant 'x' used before being initialized}}
+  x = p   // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+func testOptionalWriteGenerics2<T>(p: T) -> T? {
+  let x: T? // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = p
+  x? = p  // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+enum TestOptionalEnum {
+  case Cons(Int)
+  case Nil()
+}
+
+func testOptionalWithEnum(p: TestOptionalEnum) -> TestOptionalEnum? {
+  let x: TestOptionalEnum? // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = p
+  x? = p  // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+// Tests for optional chaining
+
+class DIOptionalTestClass {
+  var r: DIOptionalTestClass? = nil
+  var f: Int = 0;
+  let g: Int = 0;
+}
+
+func testOptionalChaining(p: DIOptionalTestClass?) {
+  p?.f = 2
+}
+
+func testOptionalChaining2(p: DIOptionalTestClass?) -> DIOptionalTestClass? {
+  let x: DIOptionalTestClass?
+  x = p
+  x?.f = 1
+  p?.r?.f = 2
+  return x
+}
+
+struct DIOptionalTestStruct {
+  var f: Int
+}
+
+func testOptionalChaining3() -> DIOptionalTestStruct? {
+  let x: DIOptionalTestStruct?  // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = DIOptionalTestStruct(f: 0)
+  x?.f = 2  // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+extension DIOptionalTestStruct {
+  public init?() {
+    self.f = 0
+  }
+}
+
+func testOptionalChaining4() -> DIOptionalTestStruct? {
+  let x: DIOptionalTestStruct?  // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = DIOptionalTestStruct()
+  x?.f = 2  // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+struct DIOptionalTestStructPair {
+  var pair: (Int, Int)
+}
+
+func test6() -> DIOptionalTestStructPair? {
+  let x: DIOptionalTestStructPair?  // expected-note {{change 'let' to 'var' to make it mutable}}
+  x = DIOptionalTestStructPair(pair: (0, 0))
+  x?.pair.0 = 1 // expected-error {{immutable value 'x' may only be initialized once}}
+  return x
+}
+
+func testOptionalChainingWithGenerics<T: DIOptionalTestProtocol>(p: T) -> T? {
+  let x: T? // expected-note {{constant defined here}}
+            // expected-note@-1 {{constant defined here}}
+            // expected-note@-2 {{constant defined here}}
+
+  // note that here assignment to 'f' is a call to the setter.
+  x?.f = 0  // expected-error {{constant 'x' used before being initialized}}
+            // expected-error@-1 {{constant 'x' passed by reference before being initialized}}
+  return x  // expected-error {{constant 'x' used before being initialized}}
+}