Merge pull request #19576 from lorentey/aged-indices

[stdlib] Set, Dictionary: Radically improved index validation

Author: lorentey

stdlib/public/SwiftShims/GlobalObjects.h

@@ -43,7 +43,10 @@ struct _SwiftEmptyArrayStorage _swiftEmptyArrayStorage;
 struct _SwiftDictionaryBodyStorage {
   __swift_intptr_t count;
   __swift_intptr_t capacity;
-  __swift_intptr_t scale;
+  __swift_int8_t scale;
+  __swift_int8_t reservedScale;
+  __swift_int16_t extra;
+  __swift_int32_t age;
   __swift_intptr_t seed;
   void *rawKeys;
   void *rawValues;
@@ -52,7 +55,10 @@ struct _SwiftDictionaryBodyStorage {
 struct _SwiftSetBodyStorage {
   __swift_intptr_t count;
   __swift_intptr_t capacity;
-  __swift_intptr_t scale;
+  __swift_int8_t scale;
+  __swift_int8_t reservedScale;
+  __swift_int16_t extra;
+  __swift_int32_t age;
   __swift_intptr_t seed;
   void *rawElements;
 };
@@ -86,6 +92,16 @@ struct _SwiftHashingParameters _swift_stdlib_Hashing_parameters;
 
 #ifdef __cplusplus
 
+static_assert(
+  sizeof(_SwiftDictionaryBodyStorage) ==
+    5 * sizeof(__swift_intptr_t) + sizeof(__swift_int64_t),
+  "_SwiftDictionaryBodyStorage has unexpected size");
+
+static_assert(
+  sizeof(_SwiftSetBodyStorage) ==
+    4 * sizeof(__swift_intptr_t) + sizeof(__swift_int64_t),
+  "_SwiftSetBodyStorage has unexpected size");
+
 static_assert(std::is_pod<_SwiftEmptyArrayStorage>::value,
               "empty array type should be POD");
 static_assert(std::is_pod<_SwiftEmptyDictionarySingleton>::value,

stdlib/public/core/Dictionary.swift

@@ -47,6 +47,7 @@
 //   | _count                                                    |
 //   | _capacity                                                 |
 //   | _scale                                                    |
+//   | _age                                                      |
 //   | _seed                                                     |
 //   | _rawKeys                                                  |
 //   | _rawValue                                                 |
@@ -149,6 +150,13 @@
 // dictionary storage are bounds-checked, this scheme never compromises memory
 // safety.
 //
+// As a safeguard against using invalid indices, Set and Dictionary maintain a
+// mutation counter in their storage header (`_age`). This counter gets bumped
+// every time an element is removed and whenever the contents are
+// rehashed. Native indices include a copy of this counter so that index
+// validation can verify it matches with current storage. This can't catch all
+// misuse, because counters may match by accident; but it does make indexing a
+// lot more reliable.
 //
 // Bridging
 // ========
@@ -1455,8 +1463,9 @@ extension Dictionary {
         return _variant.value(at: position)
       }
       _modify {
-        let index = _variant.ensureUniqueNative(preserving: position)
-        let address = _variant.asNative._values + index.offset
+        let native = _variant.ensureUniqueNative()
+        let bucket = native.validatedBucket(for: position)
+        let address = native._values + bucket.offset
         yield &address.pointee
         _fixLifetime(self)
       }
@@ -1485,6 +1494,27 @@ extension Dictionary {
     public var debugDescription: String {
       return _makeCollectionDescription(withTypeName: "Dictionary.Values")
     }
+
+    @inlinable
+    public mutating func swapAt(_ i: Index, _ j: Index) {
+      guard i != j else { return }
+      let (a, b): (_HashTable.Bucket, _HashTable.Bucket)
+      switch _variant {
+      case .native(let native):
+        a = native.validatedBucket(for: i)
+        b = native.validatedBucket(for: j)
+#if _runtime(_ObjC)
+      case .cocoa(let cocoa):
+        _variant.cocoaPath()
+        let native = _NativeDictionary<Key, Value>(cocoa)
+        a = native.validatedBucket(for: i)
+        b = native.validatedBucket(for: j)
+        _variant = .native(native)
+#endif
+      }
+      let isUnique = _variant.isUniquelyReferenced()
+      _variant.asNative.swapValuesAt(a, b, isUnique: isUnique)
+    }
   }
 }
 
@@ -1691,7 +1721,7 @@ extension Dictionary {
     @_frozen
     @usableFromInline
     internal enum _Variant {
-      case native(_NativeDictionary<Key, Value>.Index)
+      case native(_HashTable.Index)
 #if _runtime(_ObjC)
       case cocoa(_CocoaDictionary.Index)
 #endif
@@ -1708,7 +1738,7 @@ extension Dictionary {
 
     @inlinable
     @inline(__always)
-    internal init(_native index: _NativeDictionary<Key, Value>.Index) {
+    internal init(_native index: _HashTable.Index) {
       self.init(_variant: .native(index))
     }
 
@@ -1740,13 +1770,14 @@ extension Dictionary.Index {
 #endif
 
   @usableFromInline @_transparent
-  internal var _asNative: _NativeDictionary<Key, Value>.Index {
+  internal var _asNative: _HashTable.Index {
     switch _variant {
     case .native(let nativeIndex):
       return nativeIndex
 #if _runtime(_ObjC)
     case .cocoa:
-      _sanityCheckFailure("internal error: does not contain a native index")
+      _preconditionFailure(
+        "Attempting to access Dictionary elements using an invalid index")
 #endif
     }
   }
@@ -1756,7 +1787,8 @@ extension Dictionary.Index {
   internal var _asCocoa: _CocoaDictionary.Index {
     switch _variant {
     case .native:
-      _sanityCheckFailure("internal error: does not contain a Cocoa index")
+      _preconditionFailure(
+        "Attempting to access Dictionary elements using an invalid index")
     case .cocoa(let cocoaIndex):
       return cocoaIndex
     }
@@ -1811,14 +1843,14 @@ extension Dictionary.Index: Hashable {
     switch _variant {
     case .native(let nativeIndex):
       hasher.combine(0 as UInt8)
-      hasher.combine(nativeIndex.offset)
+      hasher.combine(nativeIndex.bucket.offset)
     case .cocoa(let cocoaIndex):
       _cocoaPath()
       hasher.combine(1 as UInt8)
       hasher.combine(cocoaIndex.currentKeyIndex)
     }
   #else
-    hasher.combine(_asNative.offset)
+    hasher.combine(_asNative.bucket.offset)
   #endif
   }
 }

stdlib/public/core/DictionaryBridging.swift

@@ -403,7 +403,8 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
     let unmanagedObjects = _UnmanagedAnyObjectArray(objects!)
     var bucket = _HashTable.Bucket(offset: Int(theState.extra.0))
     let endBucket = hashTable.endBucket
-    _precondition(bucket == endBucket || hashTable.isOccupied(bucket))
+    _precondition(bucket == endBucket || hashTable.isOccupied(bucket),
+      "Invalid fast enumeration state")
     var stored = 0
 
     // Only need to bridge once, so we can hoist it out of the loop.
@@ -528,7 +529,7 @@ extension _CocoaDictionary: _DictionaryBuffer {
   @inline(__always)
   func key(at index: Index) -> Key {
     _precondition(index.base.object === self.object, "Invalid index")
-    return index.allKeys[index.currentKeyIndex]
+    return index.key
   }
 
   @inlinable
@@ -607,6 +608,25 @@ extension _CocoaDictionary {
   }
 }
 
+extension _CocoaDictionary.Index {
+  @inlinable
+  @nonobjc
+  internal var key: AnyObject {
+    _precondition(currentKeyIndex < allKeys.value,
+      "Attempting to access Dictionary elements using an invalid index")
+    return allKeys[currentKeyIndex]
+  }
+
+  @usableFromInline
+  @nonobjc
+  internal var age: Int32 {
+    @_effects(releasenone)
+    get {
+      return _HashTable.age(for: base.object)
+    }
+  }
+}
+
 extension _CocoaDictionary.Index: Equatable {
   @inlinable
   internal static func == (

stdlib/public/core/DictionaryStorage.swift

@@ -16,16 +16,14 @@ import SwiftShims
 /// Enough bytes are allocated to hold the bitmap for marking valid entries,
 /// keys, and values. The data layout starts with the bitmap, followed by the
 /// keys, followed by the values.
-//
-// See the docs at the top of the file for more details on this type
-//
-// NOTE: The precise layout of this type is relied on in the runtime
-// to provide a statically allocated empty singleton.
-// See stdlib/public/stubs/GlobalObjects.cpp for details.
 @_fixed_layout // FIXME(sil-serialize-all)
 @usableFromInline
 @_objc_non_lazy_realization
 internal class _RawDictionaryStorage: __SwiftNativeNSDictionary {
+  // NOTE: The precise layout of this type is relied on in the runtime to
+  // provide a statically allocated empty singleton.  See
+  // stdlib/public/stubs/GlobalObjects.cpp for details.
+
   /// The current number of occupied entries in this dictionary.
   @usableFromInline
   @nonobjc
@@ -41,15 +39,37 @@ internal class _RawDictionaryStorage: __SwiftNativeNSDictionary {
   /// power of `scale`.
   @usableFromInline
   @nonobjc
-  internal final var _scale: Int
+  internal final var _scale: Int8
+
+  /// The scale corresponding to the highest `reserveCapacity(_:)` call so far,
+  /// or 0 if there were none. This may be used later to allow removals to
+  /// resize storage.
+  ///
+  /// FIXME: <rdar://problem/18114559> Shrink storage on deletion
+  @usableFromInline
+  @nonobjc
+  internal final var _reservedScale: Int8
+
+  // Currently unused, set to zero.
+  @nonobjc
+  internal final var _extra: Int16
+
+  /// A mutation count, enabling stricter index validation.
+  @usableFromInline
+  @nonobjc
+  internal final var _age: Int32
 
+  /// The hash seed used to hash elements in this dictionary instance.
   @usableFromInline
   internal final var _seed: Int
 
+  /// A raw pointer to the start of the tail-allocated hash buffer holding keys.
   @usableFromInline
   @nonobjc
   internal final var _rawKeys: UnsafeMutableRawPointer
 
+  /// A raw pointer to the start of the tail-allocated hash buffer holding
+  /// values.
   @usableFromInline
   @nonobjc
   internal final var _rawValues: UnsafeMutableRawPointer
@@ -272,7 +292,8 @@ final internal class _DictionaryStorage<Key: Hashable, Value>
     let unmanagedObjects = _UnmanagedAnyObjectArray(objects!)
     var bucket = _HashTable.Bucket(offset: Int(theState.extra.0))
     let endBucket = hashTable.endBucket
-    _precondition(bucket == endBucket || _hashTable.isOccupied(bucket))
+    _precondition(bucket == endBucket || hashTable.isOccupied(bucket),
+      "Invalid fast enumeration state")
     var stored = 0
     for i in 0..<count {
       if bucket == endBucket { break }
@@ -345,23 +366,43 @@ extension _DictionaryStorage {
     _sanityCheck(capacity >= original._count)
     let scale = _HashTable.scale(forCapacity: capacity)
     let rehash = (scale != original._scale)
-    let newStorage = _DictionaryStorage<Key, Value>.allocate(scale: scale)
+    let newStorage = _DictionaryStorage<Key, Value>.allocate(
+      scale: scale,
+      // Invalidate indices if we're rehashing.
+      age: rehash ? nil : original._age
+    )
     return (newStorage, rehash)
   }
 
   @usableFromInline
   @_effects(releasenone)
   static internal func allocate(capacity: Int) -> _DictionaryStorage {
     let scale = _HashTable.scale(forCapacity: capacity)
-    return allocate(scale: scale)
+    return allocate(scale: scale, age: nil)
   }
 
-  static internal func allocate(scale: Int) -> _DictionaryStorage {
+#if _runtime(_ObjC)
+  @usableFromInline
+  @_effects(releasenone)
+  static internal func convert(
+    _ cocoa: _CocoaDictionary,
+    capacity: Int
+  ) -> _DictionaryStorage {
+    let scale = _HashTable.scale(forCapacity: capacity)
+    let age = _HashTable.age(for: cocoa.object)
+    return allocate(scale: scale, age: age)
+  }
+#endif
+
+  static internal func allocate(
+    scale: Int8,
+    age: Int32?
+  ) -> _DictionaryStorage {
     // The entry count must be representable by an Int value; hence the scale's
     // peculiar upper bound.
     _sanityCheck(scale >= 0 && scale < Int.bitWidth - 1)
 
-    let bucketCount = 1 &<< scale
+    let bucketCount = (1 as Int) &<< scale
     let wordCount = _UnsafeBitset.wordCount(forCapacity: bucketCount)
     let storage = Builtin.allocWithTailElems_3(
       _DictionaryStorage<Key, Value>.self,
@@ -379,6 +420,18 @@ extension _DictionaryStorage {
     storage._count = 0
     storage._capacity = _HashTable.capacity(forScale: scale)
     storage._scale = scale
+    storage._reservedScale = 0
+    storage._extra = 0
+
+    if let age = age {
+      storage._age = age
+    } else {
+      // The default mutation count is simply a scrambled version of the storage
+      // address.
+      storage._age = Int32(
+        truncatingIfNeeded: ObjectIdentifier(storage).hashValue)
+    }
+
     storage._rawKeys = UnsafeMutableRawPointer(keysAddr)
     storage._rawValues = UnsafeMutableRawPointer(valuesAddr)
 
@@ -390,7 +443,7 @@ extension _DictionaryStorage {
     // FIXME: Use true per-instance seeding instead. Per-capacity seeding still
     // leaves hash values the same in same-sized tables, which may affect
     // operations on two tables at once. (E.g., union.)
-    storage._seed = scale
+    storage._seed = Int(scale)
 
     // Initialize hash table metadata.
     storage._hashTable.clear()