Merge pull request #18514 from gregomni/8398

[SILGen] Forming transitive capture lists, need to unique on value not value+flags

Author: slavapestov

include/swift/AST/CaptureInfo.h

@@ -64,6 +64,12 @@ class CapturedValue {
 
   bool isDynamicSelfMetadata() const { return !Value.getPointer(); }
 
+  CapturedValue mergeFlags(CapturedValue cv) {
+    assert(Value.getPointer() == cv.Value.getPointer() &&
+           "merging flags on two different value decls");
+    return CapturedValue(Value.getPointer(), getFlags() & cv.getFlags());
+  }
+
   ValueDecl *getDecl() const {
     assert(Value.getPointer() && "dynamic Self metadata capture does not "
            "have a value");

lib/SIL/TypeLowering.cpp

@@ -2038,7 +2038,7 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
 
   // Recursively collect transitive captures from captured local functions.
   llvm::DenseSet<AnyFunctionRef> visitedFunctions;
-  llvm::SetVector<CapturedValue> captures;
+  llvm::MapVector<ValueDecl*,CapturedValue> captures;
 
   // If there is a capture of 'self' with dynamic 'Self' type, it goes last so
   // that IRGen can pass dynamic 'Self' metadata.
@@ -2139,30 +2139,43 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
 
           // We're capturing a 'self' value with dynamic 'Self' type;
           // handle it specially.
-          if (!selfCapture &&
-              captureType->getClassOrBoundGenericClass()) {
-            selfCapture = capture;
+          if (captureType->getClassOrBoundGenericClass()) {
+            if (selfCapture)
+              selfCapture = selfCapture->mergeFlags(capture);
+            else
+              selfCapture = capture;
             continue;
           }
         }
 
         // Fall through to capture the storage.
       }
-      
+
       // Collect non-function captures.
-      captures.insert(capture);
+      ValueDecl *value = capture.getDecl();
+      auto existing = captures.find(value);
+      if (existing != captures.end()) {
+        existing->second = existing->second.mergeFlags(capture);
+      } else {
+        captures.insert(std::pair<ValueDecl *, CapturedValue>(value, capture));
+      }
     }
   };
   collectFunctionCaptures(fn);
 
+  SmallVector<CapturedValue, 4> resultingCaptures;
+  for (auto capturePair : captures) {
+    resultingCaptures.push_back(capturePair.second);
+  }
+
   // If we captured the dynamic 'Self' type and we have a 'self' value also,
   // add it as the final capture. Otherwise, add a fake hidden capture for
   // the dynamic 'Self' metatype.
   if (selfCapture.hasValue()) {
-    captures.insert(*selfCapture);
+    resultingCaptures.push_back(*selfCapture);
   } else if (capturesDynamicSelf) {
     selfCapture = CapturedValue::getDynamicSelfMetadata();
-    captures.insert(*selfCapture);
+    resultingCaptures.push_back(*selfCapture);
   }
 
   // Cache the uniqued set of transitive captures.
@@ -2171,7 +2184,7 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
   auto &cachedCaptures = inserted.first->second;
   cachedCaptures.setGenericParamCaptures(capturesGenericParams);
   cachedCaptures.setDynamicSelfType(capturesDynamicSelf);
-  cachedCaptures.setCaptures(Context.AllocateCopy(captures));
+  cachedCaptures.setCaptures(Context.AllocateCopy(resultingCaptures));
 
   return cachedCaptures;
 }

test/SILGen/capture-transitive.swift

@@ -0,0 +1,41 @@
+// RUN: %target-swift-emit-silgen -enable-sil-ownership %s | %FileCheck %s
+// SR-8398
+func fibonacci(_ n: Int) -> Int {
+  var cache: [Int: Int] = [:]
+
+  func recursive(_ m: Int) -> Int {
+    return cache[m] ?? {
+      // Make sure cache is only captured once in the closure
+      // CHECK: implicit closure #1 in recursive #1
+      // CHECK-LABEL: sil private [transparent] @{{.*}}9fibonacci{{.*}}9recursive{{.*}} : $@convention(thin) (Int, @guaranteed { var Dictionary<Int, Int> }) -> (Int, @error Error)
+      // CHECK: closure #1 in implicit closure #1 in recursive #1
+      // CHECK-LABEL: sil private @{{.*}}9fibonacci{{.*}}9recursive{{.*}} : $@convention(thin) (Int, @guaranteed { var Dictionary<Int, Int> }) -> Int
+      let output = m < 2 ? m : recursive(m - 1) + recursive(m - 2)
+      cache[m] = output
+      return output
+      }()
+  }
+  return recursive(n)
+}
+
+class C {
+    required init() {}
+    func f() {}
+    // Make sure that we capture dynamic self type if an explicit self isn't guaranteed
+    func returnsSelf() -> Self {
+        return { self.f(); return .init() }()
+        // CHECK-LABEL: sil private @{{.*}}returnsSelf{{.*}} : $@convention(thin) (@guaranteed C) -> @owned C
+    }
+
+    func returnsSelf1() -> Self {
+        return { [weak self] in self?.f(); return .init() }()
+        // CHECK-LABEL: sil private @{{.*}}returnsSelf{{.*}}  : $@convention(thin) (@guaranteed { var @sil_weak Optional<C> }, @thick @dynamic_self C.Type) -> @owned C
+    }
+
+    func returnsSelf2() -> Self {
+        return { [unowned self] in self.f(); return .init() }()
+        // CHECK-LABEL: sil private @{{.*}}returnsSelf{{.*}}  : $@convention(thin) (@guaranteed @sil_unowned C, @thick @dynamic_self C.Type) -> @owned C
+    }
+}
+
+