[stdlib] Use addressor for Dictionary's subscript(_:default:)

benchmark/CMakeLists.txt

@@ -57,6 +57,7 @@ set(SWIFT_BENCH_MODULES
     single-source/DictionaryGroup
     single-source/DictionaryLiteral
     single-source/DictionaryRemove
+    single-source/DictionarySubscriptDefault
     single-source/DictionarySwap
     single-source/DropFirst
     single-source/DropLast

benchmark/single-source/DictionarySubscriptDefault.swift

@@ -0,0 +1,124 @@
+//===--- DictionarySubscriptDefault.swift ---------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import TestsUtils
+
+public let DictionarySubscriptDefault = [
+  BenchmarkInfo(name: "DictionarySubscriptDefaultMutation",
+                runFunction: run_DictionarySubscriptDefaultMutation,
+                tags: [.validation, .api, .Dictionary]),
+  BenchmarkInfo(name: "DictionarySubscriptDefaultMutationArray",
+                runFunction: run_DictionarySubscriptDefaultMutationArray,
+                tags: [.validation, .api, .Dictionary]),
+  BenchmarkInfo(name: "DictionarySubscriptDefaultMutationOfObjects",
+                runFunction: run_DictionarySubscriptDefaultMutationOfObjects,
+                tags: [.validation, .api, .Dictionary]),
+  BenchmarkInfo(name: "DictionarySubscriptDefaultMutationArrayOfObjects",
+                runFunction:
+                  run_DictionarySubscriptDefaultMutationArrayOfObjects,
+                tags: [.validation, .api, .Dictionary]),
+]
+
+let count = 10_000
+let result = count / 100
+
+@inline(never)
+public func run_DictionarySubscriptDefaultMutation(_ N: Int) {
+  for _ in 1...N {
+
+    var dict = [Int: Int]()
+
+    for i in 0..<count {
+      dict[i % 100, default: 0] += 1
+    }
+
+    CheckResults(dict.count == 100)
+    CheckResults(dict[0]! == result)
+  }
+}
+
+@inline(never)
+public func run_DictionarySubscriptDefaultMutationArray(_ N: Int) {
+  for _ in 1...N {
+
+    var dict = [Int: [Int]]()
+
+    for i in 0..<count {
+      dict[i % 100, default: []].append(i)
+    }
+
+    CheckResults(dict.count == 100)
+    CheckResults(dict[0]!.count == result)
+  }
+}
+
+// Hack to workaround the fact that if we attempt to increment the Box's value
+// from the subscript, the compiler will just call the subscript's getter (and
+// therefore not insert the instance) as it's dealing with a reference type.
+// By using a mutating method in a protocol extension, the compiler is forced to
+// treat this an actual mutation, so cannot call the getter.
+protocol P {
+  associatedtype T
+  var value: T { get set }
+}
+
+extension P {
+  mutating func mutateValue(_ mutations: (inout T) -> Void) {
+    mutations(&value)
+  }
+}
+
+class Box<T : Hashable> : Hashable, P {
+  var value: T
+
+  init(_ v: T) {
+    value = v
+  }
+
+  var hashValue: Int {
+    return value.hashValue
+  }
+
+  static func ==(lhs: Box, rhs: Box) -> Bool {
+    return lhs.value == rhs.value
+  }
+}
+
+@inline(never)
+public func run_DictionarySubscriptDefaultMutationOfObjects(_ N: Int) {
+  for _ in 1...N {
+
+    var dict = [Box<Int>: Box<Int>]()
+
+    for i in 0..<count {
+      dict[Box(i % 100), default: Box(0)].mutateValue { $0 += 1 }
+    }
+
+    CheckResults(dict.count == 100)
+    CheckResults(dict[Box(0)]!.value == result)
+  }
+}
+
+@inline(never)
+public func run_DictionarySubscriptDefaultMutationArrayOfObjects(_ N: Int) {
+  for _ in 1...N {
+
+    var dict = [Box<Int>: [Box<Int>]]()
+
+    for i in 0..<count {
+      dict[Box(i % 100), default: []].append(Box(i))
+    }
+
+    CheckResults(dict.count == 100)
+    CheckResults(dict[Box(0)]!.count == result)
+  }
+}

benchmark/utils/main.swift

@@ -44,6 +44,7 @@ import DictionaryBridge
 import DictionaryGroup
 import DictionaryLiteral
 import DictionaryRemove
+import DictionarySubscriptDefault
 import DictionarySwap
 import DropFirst
 import DropLast
@@ -163,6 +164,7 @@ registerBenchmark(DictionaryBridge)
 registerBenchmark(DictionaryGroup)
 registerBenchmark(DictionaryLiteral)
 registerBenchmark(DictionaryRemove)
+registerBenchmark(DictionarySubscriptDefault)
 registerBenchmark(DictionarySwap)
 registerBenchmark(DropFirst)
 registerBenchmark(DropLast)

stdlib/public/core/HashedCollections.swift.gyb

@@ -2046,9 +2046,11 @@ public struct Dictionary<Key : Hashable, Value> :
     get {
       return _variantBuffer.maybeGet(key) ?? defaultValue()
     }
-    set(newValue) {
-      // FIXME(performance): this loads and discards the old value.
-      _variantBuffer.updateValue(newValue, forKey: key)
+    mutableAddressWithNativeOwner {
+      let (_, address) = _variantBuffer
+        .pointerToValue(forKey: key, insertingDefault: defaultValue)
+      return (address, Builtin.castToNativeObject(
+        _variantBuffer.asNative._storage))
     }
   }
 
@@ -4867,6 +4869,17 @@ internal enum _Variant${Self}Buffer<${TypeParametersDecl}> : _HashBuffer {
     }
   }
 
+  @_inlineable // FIXME(sil-serialize-all)
+  @_versioned // FIXME(sil-serialize-all)
+  internal mutating func ensureNativeBuffer() {
+#if _runtime(_ObjC)
+    if _fastPath(guaranteedNative) { return }
+    if case .cocoa(let cocoaBuffer) = self {
+      migrateDataToNativeBuffer(cocoaBuffer)
+    }
+#endif
+  }
+
 #if _runtime(_ObjC)
   @_inlineable // FIXME(sil-serialize-all)
   @_versioned // FIXME(sil-serialize-all)
@@ -5316,6 +5329,42 @@ internal enum _Variant${Self}Buffer<${TypeParametersDecl}> : _HashBuffer {
 #endif
     }
   }
+
+  @_inlineable // FIXME(sil-serialize-all)
+  @_versioned // FIXME(sil-serialize-all)
+  internal mutating func nativePointerToValue(
+    forKey key: Key, insertingDefault defaultValue: () -> Value
+  ) -> (inserted: Bool, pointer: UnsafeMutablePointer<Value>) {
+
+    var (i, found) = asNative._find(key, startBucket: asNative._bucket(key))
+    if found {
+      let pointer = nativePointerToValue(at: ._native(i))
+      return (inserted: false, pointer: pointer)
+    }
+
+    let minCapacity = NativeBuffer.minimumCapacity(
+      minimumCount: asNative.count + 1,
+      maxLoadFactorInverse: _hashContainerDefaultMaxLoadFactorInverse)
+
+    let (_, capacityChanged) = ensureUniqueNativeBuffer(minCapacity)
+
+    if capacityChanged {
+      i = asNative._find(key, startBucket: asNative._bucket(key)).pos
+    }
+
+    asNative.initializeKey(key, value: defaultValue(), at: i.offset)
+    asNative.count += 1
+    return (inserted: true, pointer: asNative.values + i.offset)
+  }
+
+  @_inlineable // FIXME(sil-serialize-all)
+  @_versioned // FIXME(sil-serialize-all)
+  internal mutating func pointerToValue(
+    forKey key: Key, insertingDefault defaultValue: () -> Value
+  ) -> (inserted: Bool, pointer: UnsafeMutablePointer<Value>) {
+    ensureNativeBuffer()
+    return nativePointerToValue(forKey: key, insertingDefault: defaultValue)
+  }
 %end
 
   @_inlineable // FIXME(sil-serialize-all)
@@ -5360,20 +5409,8 @@ internal enum _Variant${Self}Buffer<${TypeParametersDecl}> : _HashBuffer {
   internal mutating func insert(
     _ value: Value, forKey key: Key
   ) -> (inserted: Bool, memberAfterInsert: Value) {
-
-    if _fastPath(guaranteedNative) {
-      return nativeInsert(value, forKey: key)
-    }
-
-    switch self {
-    case .native:
-      return nativeInsert(value, forKey: key)
-#if _runtime(_ObjC)
-    case .cocoa(let cocoaBuffer):
-      migrateDataToNativeBuffer(cocoaBuffer)
-      return nativeInsert(value, forKey: key)
-#endif
-    }
+    ensureNativeBuffer()
+    return nativeInsert(value, forKey: key)
   }
 
 %if Self == 'Dictionary':
@@ -5477,20 +5514,8 @@ internal enum _Variant${Self}Buffer<${TypeParametersDecl}> : _HashBuffer {
     _ keysAndValues: S,
     uniquingKeysWith combine: (Value, Value) throws -> Value
   ) rethrows where S.Element == (Key, Value) {
-    if _fastPath(guaranteedNative) {
-      try nativeMerge(keysAndValues, uniquingKeysWith: combine)
-      return
-    }
-
-    switch self {
-    case .native:
-      try nativeMerge(keysAndValues, uniquingKeysWith: combine)
-#if _runtime(_ObjC)
-    case .cocoa(let cocoaStorage):
-      migrateDataToNativeBuffer(cocoaStorage)
-      try nativeMerge(keysAndValues, uniquingKeysWith: combine)
-#endif
-    }
+    ensureNativeBuffer()
+    try nativeMerge(keysAndValues, uniquingKeysWith: combine)
   }
 
   @_inlineable // FIXME(sil-serialize-all)

test/stdlib/Inputs/DictionaryKeyValueTypes.swift

@@ -25,6 +25,31 @@ func equalsUnordered<T : Comparable>(_ lhs: [T], _ rhs: [T]) -> Bool {
   return lhs.sorted().elementsEqual(rhs.sorted())
 }
 
+// A COW wrapper type that holds an Int.
+struct TestValueCOWTy {
+
+  class Base {
+    var value: Int
+    init(_ value: Int) { self.value = value }
+  }
+
+  private var base: Base
+  init(_ value: Int = 0) { self.base = Base(value) }
+
+  var value: Int {
+    get { return base.value }
+    set {
+      if !isKnownUniquelyReferenced(&base) {
+        base = Base(newValue)
+      } else {
+        base.value = newValue
+      }
+    }
+  }
+
+  var baseAddress: Int { return unsafeBitCast(base, to: Int.self) }
+}
+
 var _keyCount = _stdlib_AtomicInt(0)
 var _keySerial = _stdlib_AtomicInt(0)
 

validation-test/stdlib/Dictionary.swift

@@ -120,6 +120,14 @@ func getCOWFastDictionary() -> Dictionary<Int, Int> {
   return d
 }
 
+func getCOWFastDictionaryWithCOWValues() -> Dictionary<Int, TestValueCOWTy> {
+  var d = Dictionary<Int, TestValueCOWTy>(minimumCapacity: 10)
+  d[10] = TestValueCOWTy(1010)
+  d[20] = TestValueCOWTy(1020)
+  d[30] = TestValueCOWTy(1030)
+  return d
+}
+
 func getCOWSlowDictionary() -> Dictionary<TestKeyTy, TestValueTy> {
   var d = Dictionary<TestKeyTy, TestValueTy>(minimumCapacity: 10)
   d[TestKeyTy(10)] = TestValueTy(1010)
@@ -797,6 +805,32 @@ DictionaryTestSuite.test("COW.Fast.DefaultedSubscriptDoesNotReallocate") {
   }
 }
 
+DictionaryTestSuite.test("COW.Fast.DefaultedSubscriptDoesNotCopyValue") {
+  do {
+    var d = getCOWFastDictionaryWithCOWValues()
+    let identityValue30 = d[30]!.baseAddress
+
+    // Increment the value without having to reallocate the underlying Base
+    // instance, as uniquely referenced.
+    d[30, default: TestValueCOWTy()].value += 1
+    assert(identityValue30 == d[30]!.baseAddress)
+    assert(d[30]!.value == 1031)
+
+    let value40 = TestValueCOWTy()
+    let identityValue40 = value40.baseAddress
+
+    // Increment the value, reallocating the underlying Base, as not uniquely
+    // referenced.
+    d[40, default: value40].value += 1
+    assert(identityValue40 != d[40]!.baseAddress)
+    assert(d[40]!.value == 1)
+
+    // Keep variables alive.
+    _fixLifetime(d)
+    _fixLifetime(value40)
+  }
+}
+
 DictionaryTestSuite.test("COW.Fast.IndexForKeyDoesNotReallocate") {
   var d = getCOWFastDictionary()
   var identity1 = d._rawIdentifier()