[NO MERGE] Exclusive key path fix

include/swift/Runtime/Exclusivity.h

@@ -58,6 +58,19 @@ void swift_endAccess(ValueBuffer *buffer);
 SWIFT_RUNTIME_EXPORT
 bool _swift_disableExclusivityChecking;
 
+#ifndef NDEBUG
+
+/// Dump all accesses currently tracked by the runtime.
+///
+/// This is a debug routine that is intended to be used from the debugger and is
+/// compiled out when asserts are disabled. The intention is that it allows one
+/// to dump the access state to easily see if/when exclusivity violations will
+/// happen. This eases debugging.
+SWIFT_RUNTIME_EXPORT
+void swift_dumpTrackedAccesses();
+
+#endif
+
 } // end namespace swift
 
 #endif

include/swift/Strings.h

@@ -13,6 +13,9 @@
 #ifndef SWIFT_STRINGS_H
 #define SWIFT_STRINGS_H
 
+#include "swift/Basic/LLVM.h"
+#include "llvm/ADT/StringRef.h"
+
 namespace swift {
 
 /// The extension for serialized modules.
@@ -89,6 +92,10 @@ constexpr static const char BUILTIN_TYPE_NAME_VEC[] = "Builtin.Vec";
 constexpr static const char BUILTIN_TYPE_NAME_SILTOKEN[] = "Builtin.SILToken";
 /// The name of the Builtin type for Word
 constexpr static const char BUILTIN_TYPE_NAME_WORD[] = "Builtin.Word";
+
+constexpr static StringLiteral OPTIMIZE_SIL_PRESERVE_EXCLUSIVITY =
+    "optimize.sil.preserve_exclusivity";
+
 } // end namespace swift
 
 #endif // SWIFT_STRINGS_H

lib/IRGen/GenDecl.cpp

@@ -2310,11 +2310,18 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(SILFunction *f,
 
   LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
 
-  if (f->getInlineStrategy() == NoInline) {
+  switch (f->getInlineStrategy()) {
+  case NoInline:
     attrs = attrs.addAttribute(signature.getType()->getContext(),
                                llvm::AttributeList::FunctionIndex,
                                llvm::Attribute::NoInline);
+    break;
+  case AlwaysInline:
+    // We do not transfer AlwaysInline since we get weird test failures.
+  case InlineDefault:
+    break;
   }
+
   if (isReadOnlyFunction(f)) {
     attrs = attrs.addAttribute(signature.getType()->getContext(),
                                llvm::AttributeList::FunctionIndex,

lib/SILOptimizer/Mandatory/AccessMarkerElimination.cpp

@@ -25,6 +25,7 @@
 #include "swift/Basic/Range.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/Strings.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace swift;
@@ -167,6 +168,12 @@ struct AccessMarkerEliminationPass : SILModuleTransform {
   void run() override {
     auto &M = *getModule();
     for (auto &F : M) {
+      if (F.hasSemanticsAttr(OPTIMIZE_SIL_PRESERVE_EXCLUSIVITY)) {
+        DEBUG(llvm::dbgs() << "Skipping " << F.getName() << ". Found "
+                           << OPTIMIZE_SIL_PRESERVE_EXCLUSIVITY << " tag!\n");
+        continue;
+      }
+
       bool removedAny = AccessMarkerElimination(&F).stripMarkers();
 
       // Only invalidate analyses if we removed some markers.

lib/SILOptimizer/Mandatory/IRGenPrepare.cpp

@@ -20,12 +20,14 @@
 ///
 //===----------------------------------------------------------------------===//
 
-#include "swift/SILOptimizer/PassManager/Passes.h"
+#define DEBUG_TYPE "sil-cleanup"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
-#include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/Strings.h"
 
 using namespace swift;
 
@@ -68,10 +70,20 @@ namespace {
 
 class IRGenPrepare : public SILFunctionTransform {
   void run() override {
-    bool shouldInvalidate = cleanFunction(*getFunction());
-    if (!shouldInvalidate)
-      return;
-    invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
+    SILFunction *F = getFunction();
+
+    if (F->hasSemanticsAttr(OPTIMIZE_SIL_PRESERVE_EXCLUSIVITY)) {
+      assert(F->getInlineStrategy() == NoInline &&
+             "All OPTIMIZE_SIL_PRESERVE_EXCLUSIVITY functions should have "
+             "no-inline");
+      // We always want to inline these at the llvm level.
+      F->setInlineStrategy(AlwaysInline);
+    }
+
+    bool shouldInvalidate = cleanFunction(*F);
+
+    if (shouldInvalidate)
+      invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
   }
 };
 

stdlib/public/core/KeyPath.swift

@@ -625,23 +625,89 @@ internal enum KeyPathComponent: Hashable {
   }
 }
 
+
+@usableFromInline
+@inline(never)
+@_semantics("optimize.sil.preserve_exclusivity")
+func beginAccessHelper<T>(_ address: Builtin.RawPointer, _ accessRecordPtr: Builtin.RawPointer, _ type: T.Type) {
+  print("Begin: \(address) - \(accessRecordPtr)")
+  Builtin.beginUnpairedModifyAccess(address, accessRecordPtr, type)
+}
+
+@usableFromInline
+@inline(never)
+@_semantics("optimize.sil.preserve_exclusivity")
+func endAccessHelper(_ accessRecordPtr: Builtin.RawPointer) {
+  print("End: \(accessRecordPtr)")
+  Builtin.endUnpairedAccess(accessRecordPtr)
+}
+
 // A class that maintains ownership of another object while a mutable projection
-// into it is underway.
+// into it is underway. The lifetime of the instance of this class is also used
+// to begin and end exclusive 'modify' access to the projected address.
 @_fixed_layout // FIXME(sil-serialize-all)
 @usableFromInline // FIXME(sil-serialize-all)
-internal final class ClassHolder {
+internal final class ClassHolder<ProjectionType> {
+
+  /// The type of the scratch record passed to the runtime to record
+  /// accesses to guarantee exlcusive access.
+  internal typealias AccessRecord = Builtin.UnsafeValueBuffer
+
   @usableFromInline // FIXME(sil-serialize-all)
-  internal let previous: AnyObject?
+  internal var previous: AnyObject?
   @usableFromInline // FIXME(sil-serialize-all)
-  internal let instance: AnyObject
+  internal var instance: AnyObject
 
   @inlinable // FIXME(sil-serialize-all)
-  internal init(previous: AnyObject?, instance: AnyObject) {
-    self.previous = previous
-    self.instance = instance
+  internal init() {
+    _sanityCheckFailure("use _create(...)")
   }
+
   @inlinable // FIXME(sil-serialize-all)
-  deinit {}
+  internal final class func _create(
+      previous: AnyObject?,
+      instance: AnyObject,
+      accessingAddress address: UnsafeRawPointer,
+      type: ProjectionType.Type
+  ) -> ClassHolder {
+
+    // Tail allocate the UnsafeValueBuffer used as the AccessRecord.
+    // This avoids a second heap allocation since there is no source-level way to
+    // initialize a Builtin.UnsafeValueBuffer type and thus we cannot have a
+    // stored property of that type.
+    let holder: ClassHolder = Builtin.allocWithTailElems_1(self,
+                                                          1._builtinWordValue,
+                                                          AccessRecord.self)
+
+    // Initialize the ClassHolder's instance variables. This is done via
+    // withUnsafeMutablePointer(to:) because the instance was just allocated with
+    // allocWithTailElems_1 and so we need to make sure to use an initialization
+    // rather than an assignment.
+    withUnsafeMutablePointer(to: &holder.previous) {
+      $0.initialize(to: previous)
+    }
+
+    withUnsafeMutablePointer(to: &holder.instance) {
+      $0.initialize(to: instance)
+    }
+
+    let accessRecordPtr = Builtin.projectTailElems(holder, AccessRecord.self)
+
+    // Begin a 'modify' access to the address. This access is ended in
+    // ClassHolder's deinitializer.
+    beginAccessHelper(address._rawValue, accessRecordPtr, type)
+
+    return holder
+  }
+
+  @inlinable // FIXME(sil-serialize-all)
+  @usableFromInline // FIXME(sil-serialize-all)
+  deinit {
+    let accessRecordPtr = Builtin.projectTailElems(self, AccessRecord.self)
+
+    // Ends the access begun in _create().
+    endAccessHelper(accessRecordPtr)
+  }
 }
 
 // A class that triggers writeback to a pointer when destroyed.
@@ -1185,7 +1251,15 @@ internal struct RawKeyPathComponent {
       let baseObj = unsafeBitCast(base, to: AnyObject.self)
       let basePtr = UnsafeRawPointer(Builtin.bridgeToRawPointer(baseObj))
       defer { _fixLifetime(baseObj) }
-      return .continue(basePtr.advanced(by: offset)
+
+      let offsetAddress = basePtr.advanced(by: offset)
+
+      // Perform an instaneous record access on the address in order to
+      // ensure that the read will not conflict with an already in-progress
+      // 'modify' access.
+      Builtin.performInstantaneousReadAccess(offsetAddress._rawValue,
+                                             NewValue.self)
+      return .continue(offsetAddress
         .assumingMemoryBound(to: NewValue.self)
         .pointee)
 
@@ -1244,12 +1318,16 @@ internal struct RawKeyPathComponent {
       // AnyObject memory can alias any class reference memory, so we can
       // assume type here
       let object = base.assumingMemoryBound(to: AnyObject.self).pointee
-      // The base ought to be kept alive for the duration of the derived access
-      keepAlive = keepAlive == nil
-        ? object
-        : ClassHolder(previous: keepAlive, instance: object)
-      return UnsafeRawPointer(Builtin.bridgeToRawPointer(object))
+      let offsetAddress = UnsafeRawPointer(Builtin.bridgeToRawPointer(object))
             .advanced(by: offset)
+
+      // Keep the  base alive for the duration of the derived access and also
+      // enforce exclusive access to the address.
+      keepAlive = ClassHolder._create(previous: keepAlive, instance: object,
+                                      accessingAddress: offsetAddress,
+                                      type: NewValue.self)
+
+      return offsetAddress
 
     case .mutatingGetSet(id: _, get: let rawGet, set: let rawSet,
                          argument: let argument):

stdlib/public/runtime/Exclusivity.cpp

@@ -260,6 +260,16 @@ class AccessSet {
 
     swift_runtime_unreachable("access not found in set");
   }
+
+#ifndef NDEBUG
+  /// Only available with asserts. Intended to be used with
+  /// swift_dumpTrackedAccess().
+  void forEach(std::function<void (Access *)> action) {
+    for (auto *iter = Head; iter != nullptr; iter = iter->getNext()) {
+      action(iter);
+    }
+  }
+#endif
 };
 
 } // end anonymous namespace
@@ -348,3 +358,17 @@ void swift::swift_endAccess(ValueBuffer *buffer) {
 
   getAccessSet().remove(access);
 }
+
+#ifndef NDEBUG
+
+// Dump the accesses that are currently being tracked by the runtime.
+//
+// This is only intended to be used in the debugger.
+void swift::swift_dumpTrackedAccesses() {
+  getAccessSet().forEach([](Access *a) {
+      fprintf(stderr, "Access. Pointer: %p. PC: %p. AccessAction: %s",
+              a->Pointer, a->PC, getAccessName(a->getAccessAction()));
+  });
+}
+
+#endif

test/IRGen/preserve_exclusivity.swift

@@ -0,0 +1,102 @@
+// RUN: %target-swift-frontend -parse-stdlib -parse-stdlib -emit-ir -Onone %s | %FileCheck --check-prefix=IR-Onone %s
+
+// We check separately that:
+//
+// 1. We properly emit our special functions as inline always.
+// 2. end-to-end we inline the access markers.
+
+// RUN: %target-swift-frontend -parse-stdlib -Xllvm -sil-disable-pass=FunctionSignatureOpts -Xllvm -sil-disable-pass=GenericSpecializer -parse-stdlib -emit-ir -O -disable-llvm-optzns %s | %FileCheck --check-prefix=IR-Osil %s
+// RUN: %target-swift-frontend -parse-stdlib -Xllvm -sil-disable-pass=FunctionSignatureOpts -Xllvm -sil-disable-pass=GenericSpecializer -parse-stdlib -emit-ir -O %s | %FileCheck --check-prefix=IR-Ollvm %s
+
+@_silgen_name("marker1")
+func marker1() -> ()
+
+@_silgen_name("marker2")
+func marker2() -> ()
+
+@_silgen_name("marker3")
+func marker3() -> ()
+
+@_silgen_name("marker4")
+func marker4() -> ()
+
+// IR-Onone: define swiftcc void @"$S20preserve_exclusivity10beginNoOptyyBp_BpxmtlF"(i8*, i8*, %swift.type*, %swift.type* %T1)
+// IR-Onone: call void @swift_beginAccess
+// IR-Onone-NEXT: ret void
+
+// IR-Osil: define swiftcc void @"$S20preserve_exclusivity10beginNoOptyyBp_BpxmtlF"(i8*, i8*, %swift.type*, %swift.type* %T1) [[ATTR:#[0-9][0-9]*]] {
+// IR-Osil:   call void @swift_beginAccess
+// IR-Osil-NEXT: ret void
+
+@inline(never)
+@_semantics("optimize.sil.preserve_exclusivity")
+public func beginNoOpt<T1>(_ address: Builtin.RawPointer, _ scratch: Builtin.RawPointer, _ ty1: T1.Type) {
+  marker1()
+  Builtin.beginUnpairedModifyAccess(address, scratch, ty1);
+}
+
+// IR-Onone: define swiftcc void @"$S20preserve_exclusivity8endNoOptyyBpF"(i8*)
+// IR-Onone: call void @swift_endAccess
+// IR-Onone-NEXT: ret void
+
+// IR-Osil: define swiftcc void @"$S20preserve_exclusivity8endNoOptyyBpF"(i8*) [[ATTR]]
+// IR-Osil:   call void @swift_endAccess
+// IR-Osil-NEXT: ret void
+@inline(never)
+@_semantics("optimize.sil.preserve_exclusivity")
+public func endNoOpt(_ address: Builtin.RawPointer) {
+  marker2()
+  Builtin.endUnpairedAccess(address)
+}
+
+class Klass {}
+
+// Make sure testNoOpt properly inlines in our functions.
+//
+// IR-Ollvm: define swiftcc void @"$S20preserve_exclusivity9testNoOptyyBpF"(i8*)
+// IR-Ollvm: call swiftcc void @marker1
+// IR-Ollvm: call void @swift_beginAccess
+// IR-Ollvm: call swiftcc void @marker2
+// IR-Ollvm: call void @swift_endAccess
+// IR-Ollvm-NEXT: ret void
+public func testNoOpt(_ k1: Builtin.RawPointer) {
+  beginNoOpt(k1, k1, Builtin.RawPointer.self)
+  endNoOpt(k1)
+}
+
+// IR-Onone: define swiftcc void @"$S20preserve_exclusivity8beginOptyyBp_BpxmtlF"(i8*, i8*, %swift.type*, %swift.type* %T1)
+// IR-Onone: call void @swift_beginAccess
+// IR-Onone-NEXT: ret void
+
+// IR-Osil: define swiftcc void @"$S20preserve_exclusivity8beginOptyyBp_BpxmtlF"(i8*, i8*, %swift.type*, %swift.type* %T1)
+// IR-Osil-NEXT: entry
+// IR-Osil-NEXT: call swiftcc void @marker3
+// IR-Osil-NEXT: ret void
+
+@inline(never)
+public func beginOpt<T1>(_ address: Builtin.RawPointer, _ scratch: Builtin.RawPointer, _ ty1: T1.Type) {
+  marker3()
+  Builtin.beginUnpairedModifyAccess(address, scratch, ty1);
+}
+
+// IR-Onone: define swiftcc void @"$S20preserve_exclusivity6endOptyyBpF"(i8*)
+// IR-Onone: call void @swift_endAccess
+// IR-Onone-NEXT: ret void
+
+// IR-Osil: define swiftcc void @"$S20preserve_exclusivity6endOptyyBpF"(i8*)
+// IR-Osil-NEXT: entry
+// IR-Osil-NEXT: call swiftcc void @marker4
+// IR-Osil-NEXT: ret void
+
+@inline(never)
+public func endOpt(_ address: Builtin.RawPointer) {
+  marker4()
+  Builtin.endUnpairedAccess(address)
+}
+
+public func testOpt(_ k1: Builtin.RawPointer) {
+  beginOpt(k1, k1, Builtin.RawPointer.self)
+  endOpt(k1)
+}
+
+// IR-Osil: attributes [[ATTR]] = { alwaysinline