Merge pull request swiftlang#35276 from eeckstein/di-implicit-closures

DefiniteInitialization: correctly handle implicit closures.

Author: eeckstein

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -475,9 +475,14 @@ namespace {
 /// Gathers information about a specific address and its uses to determine
 /// definite initialization.
 class ElementUseCollector {
+public:
+  typedef SmallPtrSet<SILFunction *, 8> FunctionSet;
+
+private:
   SILModule &Module;
   const DIMemoryObjectInfo &TheMemory;
   DIElementUseInfo &UseInfo;
+  FunctionSet &VisitedClosures;
 
   /// IsSelfOfNonDelegatingInitializer - This is true if we're looking at the
   /// top level of a 'self' variable in a non-delegating init method.
@@ -494,12 +499,15 @@ class ElementUseCollector {
 
 public:
   ElementUseCollector(const DIMemoryObjectInfo &TheMemory,
-                      DIElementUseInfo &UseInfo)
-      : Module(TheMemory.getModule()), TheMemory(TheMemory), UseInfo(UseInfo) {}
+                      DIElementUseInfo &UseInfo,
+                      FunctionSet &visitedClosures)
+      : Module(TheMemory.getModule()), TheMemory(TheMemory), UseInfo(UseInfo),
+        VisitedClosures(visitedClosures)
+  {}
 
   /// This is the main entry point for the use walker.  It collects uses from
   /// the address and the refcount result of the allocation.
-  void collectFrom() {
+  void collectFrom(SILValue V, bool collectDestroysOfContainer) {
     IsSelfOfNonDelegatingInitializer = TheMemory.isNonDelegatingInit();
 
     // If this is a delegating initializer, collect uses specially.
@@ -508,12 +516,16 @@ class ElementUseCollector {
       assert(!TheMemory.isDerivedClassSelfOnly() &&
              "Should have been handled outside of here");
       // If this is a class pointer, we need to look through ref_element_addrs.
-      collectClassSelfUses();
+      collectClassSelfUses(V);
       return;
     }
 
-    collectUses(TheMemory.getUninitializedValue(), 0);
-    gatherDestroysOfContainer(TheMemory, UseInfo);
+    collectUses(V, 0);
+    if (collectDestroysOfContainer) {
+      assert(V == TheMemory.getUninitializedValue() &&
+             "can only gather destroys of root value");
+      gatherDestroysOfContainer(TheMemory, UseInfo);
+    }
   }
 
   void trackUse(DIMemoryUse Use) { UseInfo.trackUse(Use); }
@@ -525,7 +537,9 @@ class ElementUseCollector {
 
 private:
   void collectUses(SILValue Pointer, unsigned BaseEltNo);
-  void collectClassSelfUses();
+  bool addClosureElementUses(PartialApplyInst *pai, Operand *argUse);
+
+  void collectClassSelfUses(SILValue ClassPointer);
   void collectClassSelfUses(SILValue ClassPointer, SILType MemorySILType,
                             llvm::SmallDenseMap<VarDecl *, unsigned> &EN);
 
@@ -911,9 +925,15 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
       continue;
     }
 
+    if (User->isDebugInstruction())
+      continue;
+
     if (auto *PAI = dyn_cast<PartialApplyInst>(User)) {
       if (onlyUsedByAssignByWrapper(PAI))
         continue;
+        
+      if (BaseEltNo == 0 && addClosureElementUses(PAI, Op))
+        continue;
     }
 
     // Sanitizer instrumentation is not user visible, so it should not
@@ -926,9 +946,85 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
   }
 }
 
+/// Add all used elements of an implicit closure, which is capturing 'self'.
+///
+/// We want to correctly handle implicit closures in initializers, e.g. with
+/// boolean operators:
+/// \code
+///   init() {
+///     bool_member1 = false
+///     bool_member2 = false || bool_member1 // implicit closure
+///   }
+/// \endcode
+///
+/// The implicit closure ('bool_member1' at the RHS of the || operator) captures
+/// the whole self, but only uses 'bool_member1'.
+/// If we would add the whole captured 'self' as use, we would get a
+/// "'self.bool_member2' not initialized" error at the partial_apply.
+/// Therefore we look into the body of the closure and only add the actually
+/// used members.
+bool ElementUseCollector::addClosureElementUses(PartialApplyInst *pai,
+                                                Operand *argUse) {
+  SILFunction *callee = pai->getReferencedFunctionOrNull();
+  if (!callee)
+    return false;
+
+  // Implicit closures are "transparent", which means they are always inlined.
+  // It would probably also work to handle non-transparent closures (e.g.
+  // explicit closures). But if the closure is not inlined we could end up
+  // passing a partially initialized self to the closure function. Although it
+  // would probably not cause any real problems, an `@in_guaranteed` argument
+  // (the captured 'self') is assumed to be fully initialized in SIL.
+  if (!callee->isTransparent())
+    return false;
+
+  // Implicit closures are only partial-applied once and there cannot be a
+  // recursive cycle of implicit closures.
+  // Nevertheless such a scenario is theoretically possible in SIL. To be on the
+  // safe side, check for cycles.
+  if (!VisitedClosures.insert(callee).second)
+    return false;
+
+  unsigned argIndex = ApplySite(pai).getCalleeArgIndex(*argUse);
+  SILArgument *arg = callee->getArgument(argIndex);
+  
+  // Bail if arg is not the original 'self' object, but e.g. a projected member.
+  assert(TheMemory.getType().isObject());
+  if (arg->getType().getObjectType() != TheMemory.getType())
+    return false;
+
+  DIElementUseInfo ArgUseInfo;
+  ElementUseCollector collector(TheMemory, ArgUseInfo, VisitedClosures);
+  collector.collectFrom(arg, /*collectDestroysOfContainer*/ false);
+
+  if (!ArgUseInfo.Releases.empty() || !ArgUseInfo.StoresToSelf.empty())
+    return false;
+    
+  for (const DIMemoryUse &use : ArgUseInfo.Uses) {
+    // Only handle loads and escapes. Implicit closures will not have stores or
+    // store-like uses, anyway.
+    // Also, as we don't do a flow-sensitive analysis of the callee, we cannot
+    // handle stores, because we don't know if they are unconditional or not.
+    switch (use.Kind) {
+      case DIUseKind::Load:
+      case DIUseKind::Escape:
+      case DIUseKind::InOutArgument:
+        break;
+      default:
+        return false;
+    }
+  }
+
+  // Track all uses of the closure.
+  for (const DIMemoryUse &use : ArgUseInfo.Uses) {
+    trackUse(DIMemoryUse(pai, use.Kind, use.FirstElement, use.NumElements));
+  }
+  return true;
+}
+
 /// collectClassSelfUses - Collect all the uses of a 'self' pointer in a class
 /// constructor.  The memory object has class type.
-void ElementUseCollector::collectClassSelfUses() {
+void ElementUseCollector::collectClassSelfUses(SILValue ClassPointer) {
   assert(IsSelfOfNonDelegatingInitializer &&
          TheMemory.getASTType()->getClassOrBoundGenericClass() != nullptr);
 
@@ -952,8 +1048,7 @@ void ElementUseCollector::collectClassSelfUses() {
   // If we are looking at the init method for a root class, just walk the
   // MUI use-def chain directly to find our uses.
   if (TheMemory.isRootSelf()) {
-    collectClassSelfUses(TheMemory.getUninitializedValue(), TheMemory.getType(),
-                         EltNumbering);
+    collectClassSelfUses(ClassPointer, TheMemory.getType(), EltNumbering);
     return;
   }
 
@@ -1307,11 +1402,18 @@ void ElementUseCollector::collectClassSelfUses(
     if (isUninitializedMetatypeInst(User))
       continue;
 
+    if (User->isDebugInstruction())
+      continue;
+ 
     // If this is a partial application of self, then this is an escape point
     // for it.
     if (auto *PAI = dyn_cast<PartialApplyInst>(User)) {
       if (onlyUsedByAssignByWrapper(PAI))
         continue;
+
+      if (addClosureElementUses(PAI, Op))
+        continue;
+
       Kind = DIUseKind::Escape;
     }
 
@@ -1709,5 +1811,8 @@ void swift::ownership::collectDIElementUsesFrom(
     return;
   }
 
-  ElementUseCollector(MemoryInfo, UseInfo).collectFrom();
+  ElementUseCollector::FunctionSet VisitedClosures;
+  ElementUseCollector collector(MemoryInfo, UseInfo, VisitedClosures);
+  collector.collectFrom(MemoryInfo.getUninitializedValue(),
+                        /*collectDestroysOfContainer*/ true);
 }

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -98,11 +98,8 @@ class DIMemoryObjectInfo {
   /// instruction. For alloc_box though it returns the project_box associated
   /// with the memory info.
   SingleValueInstruction *getUninitializedValue() const {
-    if (auto *mui = dyn_cast<MarkUninitializedInst>(MemoryInst)) {
-      if (auto *pbi = mui->getSingleUserOfType<ProjectBoxInst>()) {
-        return pbi;
-      }
-    }
+    if (auto *pbi = MemoryInst->getSingleUserOfType<ProjectBoxInst>())
+      return pbi;
     return MemoryInst;
   }
 

test/SILOptimizer/definite_init_closures.sil

@@ -0,0 +1,80 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -definite-init
+
+// Test that implicit closures don't count as use of the whole self and only
+// actual uses inside the closures are taken into account.
+
+// Check that no errors are generated. No FileCheck needed.
+
+sil_stage raw
+
+import Builtin
+import Swift
+import SwiftShims
+
+struct SimpleStruct {
+  @_hasStorage let x: Bool
+  @_hasStorage let y: Bool
+  init()
+}
+
+class SimpleClass {
+  @_hasStorage final let x: Bool
+  @_hasStorage final let y: Bool
+  init()
+}
+
+sil [ossa] @$s4test1SVACycfC : $@convention(method) (Bool) -> SimpleStruct {
+bb0(%0 : $Bool):
+  %1 = alloc_stack $SimpleStruct, var, name "self"
+  %2 = mark_uninitialized [rootself] %1 : $*SimpleStruct
+  %7 = begin_access [modify] [static] %2 : $*SimpleStruct
+  %8 = struct_element_addr %7 : $*SimpleStruct, #SimpleStruct.x
+  assign %0 to %8 : $*Bool
+  end_access %7 : $*SimpleStruct
+  %16 = function_ref @implicit_closure_struct : $@convention(thin) (@inout_aliasable SimpleStruct) -> (Bool, @error Error)
+  %17 = partial_apply [callee_guaranteed] %16(%2) : $@convention(thin) (@inout_aliasable SimpleStruct) -> (Bool, @error Error)
+  destroy_value %17 : $@callee_guaranteed () -> (Bool, @error Error)
+  %23 = begin_access [modify] [static] %2 : $*SimpleStruct
+  %24 = struct_element_addr %23 : $*SimpleStruct, #SimpleStruct.y
+  assign %0 to %24 : $*Bool
+  end_access %23 : $*SimpleStruct
+  %27 = load [trivial] %2 : $*SimpleStruct
+  dealloc_stack %1 : $*SimpleStruct
+  return %27 : $SimpleStruct
+}
+
+sil private [transparent] [ossa] @implicit_closure_struct : $@convention(thin) (@inout_aliasable SimpleStruct) -> (Bool, @error Error) {
+bb0(%0 : $*SimpleStruct):
+  debug_value_addr %0 : $*SimpleStruct, var, name "self", argno 2
+  %3 = begin_access [read] [static] %0 : $*SimpleStruct
+  %4 = struct_element_addr %3 : $*SimpleStruct, #SimpleStruct.x
+  %5 = load [trivial] %4 : $*Bool
+  end_access %3 : $*SimpleStruct
+  return %5 : $Bool
+}
+
+sil [ossa] @$s4test11SimpleClassC1bACSb_tcfc : $@convention(method) (Bool, @owned SimpleClass) -> @owned SimpleClass {
+bb0(%0 : $Bool, %1 : @owned $SimpleClass):
+  %4 = mark_uninitialized [rootself] %1 : $SimpleClass
+  %5 = begin_borrow %4 : $SimpleClass
+  %6 = ref_element_addr %5 : $SimpleClass, #SimpleClass.x
+  assign %0 to %6 : $*Bool
+  end_borrow %5 : $SimpleClass
+  %9 = begin_borrow %4 : $SimpleClass
+  %11 = function_ref @implicit_closure_class : $@convention(thin) (@guaranteed SimpleClass) -> (Bool, @error Error)
+  %12 = copy_value %4 : $SimpleClass
+  %13 = partial_apply [callee_guaranteed] %11(%12) : $@convention(thin) (@guaranteed SimpleClass) -> (Bool, @error Error)
+  destroy_value %13 : $@callee_guaranteed () -> (Bool, @error Error)
+  %19 = ref_element_addr %9 : $SimpleClass, #SimpleClass.y
+  assign %0 to %19 : $*Bool
+  end_borrow %9 : $SimpleClass
+  return %4 : $SimpleClass
+}
+
+sil private [transparent] [ossa] @implicit_closure_class : $@convention(thin) (@guaranteed SimpleClass) -> (Bool, @error Error) {
+bb0(%0 : @guaranteed $SimpleClass):
+  debug_value %0 : $SimpleClass, let, name "self", argno 2
+  %3 = ref_element_addr %0 : $SimpleClass, #SimpleClass.x
+  %4 = load [trivial] %3 : $*Bool
+  return %4 : $Bool
+}

test/SILOptimizer/definite_init_closures.swift

@@ -0,0 +1,72 @@
+// RUN: %target-swift-frontend -emit-sil %s -o /dev/null
+
+// Test boolean operators with implicit closures
+
+struct Simple {
+  let x: Bool
+  let y: Bool
+
+  init() {
+    x = false
+    y = false || x
+  }
+}
+
+struct NestedClosures {
+  let x: Bool
+  let y: Bool
+  let z: Bool
+
+  init(a: Bool) {
+    x = false
+    y = false
+    z = false || (y || (x || a))
+  }
+
+  init(b: Bool) {
+    x = false
+    y = false
+    // With explicit self
+    z = false || (self.y || (self.x || b))
+  }
+}
+
+class SimpleClass {
+  let x: Bool
+  let y: Bool
+
+  init() {
+    x = false
+    y = false || x
+  }
+}
+
+func forward(_ b: inout Bool) -> Bool {
+  return b
+}
+
+struct InoutUse {
+  var x: Bool
+  var y: Bool
+
+  init() {
+    x = false
+    y = false || forward(&x)
+  }
+}
+
+protocol P {
+  var b: Bool { get }
+}
+
+struct Generic<T : P> {
+  let x: T
+  let y: Bool
+
+  init(_ t: T) {
+    x = t
+    y = false || x.b
+  }
+}
+
+