feat: Add schema extension support, member JSON name extension

Sources/Smithy/Schema/Schema.swift

@@ -72,6 +72,8 @@ public final class Schema: Sendable {
     /// Will be `nil` for any schema other than a member.
     public let containerType: ShapeType?
 
+    private let _extensions = UniquelyIndexedMutableCollection([])
+
     /// Creates a new Schema using the passed parameters.
     ///
     /// No validation is performed on the parameters since calls to this initializer
@@ -129,4 +131,33 @@ public final class Schema: Sendable {
     public var value: Schema {
         members[1] // `value` will be the second member in a map or document schema, after `key`
     }
+
+    /// Gets the requested type of schema extension for this schema.
+    /// - Parameter type: The type of schema extension to be retrieved & returned.
+    /// - Returns: A schema extension of the requested type, or `nil` if none exists.
+    public func getExtension<Extension: SchemaExtension>(_ type: Extension.Type) -> Extension? {
+        _extensions.get(Extension.self)
+    }
+
+    /// Stores the passed schema extension.
+    ///
+    /// If multiple callers attempt to create a schema extension from multiple threads at the same time, each will create & store the value.
+    /// This may result in extra work performed, but the type is thread-safe and the time spent blocking any calling thread is minimized.
+    /// - Parameter value: The schema extension to be stored.
+    public func setExtension<Extension: SchemaExtension>(_ value: Extension) {
+        _extensions.set(value)
+    }
+
+    /// Gets the requested type of schema extension for this schema, creating & storing it first if it does not exist.
+    ///
+    /// Creation & storage is not performed atomically; see ``setExtension(_:)`` for the implications when
+    /// this method is called from multiple threads at the same time.
+    /// - Parameter type: The type of schema extension to be retrieved & returned.
+    /// - Returns: A schema extension of the requested type.
+    public func getOrCreateExtension<Extension: SchemaExtension>(_ type: Extension.Type) throws -> Extension {
+        if let storedExtension = getExtension(Extension.self) { return storedExtension }
+        let newExtension = try Extension(schema: self)
+        setExtension(newExtension)
+        return newExtension
+    }
 }

Sources/Smithy/Schema/SchemaExtension.swift

@@ -0,0 +1,25 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+/// A schema extension is an object that contains data which a serializer or deserializer wishes to
+/// reuse when processing the same schema in the future.
+///
+/// Schema extensions are stored on the schema itself the first time a schema is accessed, and
+/// retrieved from the schema during subsequent accesses.
+@_spi(SchemaBasedSerde)
+public protocol SchemaExtension: AnyObject, UniquelyIndexedByType {
+
+    /// Creates the schema extension by deriving its contents from the passed schema.
+    ///
+    /// This initializer allows ``Schema/getOrCreateExtension(_:)`` to create an extension on demand.
+    /// - Parameter schema: The schema that this extension will be stored on & derives its data from.
+    init(schema: Schema) throws
+}
+
+/// Call `getNextIndex()` on this counter to generate a unique index for each type of schema extension.
+@_spi(SchemaBasedSerde)
+public let schemaExtensionUniqueIndexCounter = UniqueIndexCounter()

Sources/Smithy/UniquelyIndexed/UniquelyIndexedMutableCollection.swift

@@ -0,0 +1,95 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import class Foundation.NSLock
+
+/// A mutable collection of uniquely indexed values that provides O(1) access to elements.
+///
+/// Elements are stored in a sparse array of pointers to elements, at their own unique index in the sparse array.
+/// A lock is used to enforce exclusive access to `_storage`.  The type is designed to lock for as little time
+/// as possible, so as to not cause problems in Swift concurrency.
+///
+/// A non-recursive lock is used because no method on this type acquires the lock while already holding it;
+/// it is roughly twice as fast as a recursive lock, and access to this type is on the serialization hot path.
+final class UniquelyIndexedMutableCollection: @unchecked Sendable {
+    private var _storage: [(any UniquelyIndexedByType)?]
+
+    private let lock = NSLock()
+
+    /// Creates a uniquely indexed collection from an array of uniquely indexed instances.
+    /// - Parameter collection: The array of instances to be stored.
+    init(_ collection: [any UniquelyIndexedByType]) {
+        let highestIndex = collection.map { $0.uniqueIndex }.max() ?? -1
+        var storage: [(any UniquelyIndexedByType)?] = Array(repeating: nil, count: highestIndex + 1)
+        collection.forEach { storage[$0.uniqueIndex] = $0 }
+        self._storage = storage
+    }
+
+    /// Gets the element of the collection that matches the passed type.
+    /// - Parameter _: The type of the element to be returned
+    /// - Returns: The element of the requested type, or `nil` if there is no element of that type.
+    func get<T: UniquelyIndexedByType>(_ _: T.Type) -> T? {
+        lock.lock()
+        defer { lock.unlock() }
+        guard T.uniqueIndex < _storage.count else { return nil }
+        return _storage[T.uniqueIndex] as? T
+    }
+
+    /// Sets the passed value as the stored value for that type, replacing any previously stored value.
+    ///
+    /// Use ``clear(_:)`` to remove a stored value.
+    /// - Parameter value: The element to be stored in the collection.
+    func set<T: UniquelyIndexedByType>(_ value: T) {
+        lock.lock()
+        defer { lock.unlock() }
+        if T.uniqueIndex >= _storage.count {
+            let additionalSlots = T.uniqueIndex - _storage.count + 1
+            _storage.append(contentsOf: Array(repeating: nil, count: additionalSlots))
+        }
+        _storage[T.uniqueIndex] = value
+    }
+
+    // The members below round out the collection's API but currently have no callers in production
+    // code; they are exercised by tests only.  Suppress the analyzer's unused-declaration rule rather
+    // than delete them, so the type remains a complete, symmetric collection.
+    // swiftlint:disable unused_declaration
+
+    /// Sets the stored value to `nil` for the passed type.
+    ///
+    /// Capacity in the storage is not added or reduced by this method.
+    /// - Parameter type: The type of the element to be set to `nil`.
+    func clear<T: UniquelyIndexedByType>(_ type: T.Type) {
+        lock.lock()
+        defer { lock.unlock() }
+        if T.uniqueIndex < _storage.count {
+            _storage[T.uniqueIndex] = nil
+        }
+    }
+
+    /// The number of elements in the collection.
+    var count: Int {
+        lock.lock()
+        defer { lock.unlock() }
+        return _storage.reduce(0) { $0 + ($1 != nil ? 1 : 0) }
+    }
+
+    /// Whether the collection has no elements.
+    var isEmpty: Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return !_storage.contains { $0 != nil }
+    }
+
+    /// All of the elements in the collection, returned in unique index order.
+    var allElements: [any UniquelyIndexedByType] {
+        lock.lock()
+        defer { lock.unlock() }
+        return _storage.compactMap { $0 }
+    }
+
+    // swiftlint:enable unused_declaration
+}

Sources/SmithyJSON/JSONNameExtension.swift

@@ -0,0 +1,36 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@_spi(SchemaBasedSerde)
+import class Smithy.JSONNameTrait
+@_spi(SchemaBasedSerde)
+import class Smithy.Schema
+@_spi(SchemaBasedSerde)
+import protocol Smithy.SchemaExtension
+@_spi(SchemaBasedSerde)
+import var Smithy.schemaExtensionUniqueIndexCounter
+
+/// Stores the UTF-8 bytes for the JSON string representation of this member's JSON name.  If there
+/// is a JSON name trait on this member, it is used in place of the member component of the shape ID.
+///
+/// Bytes may be reused when this member is sent again, increasing performance by not requiring
+/// recalculation of the JSON bytes.  The tradeoff is the extra memory consumed to store the bytes
+/// in this schema extension.
+///
+/// Bytes include leading & trailing double-quotes, and all characters requiring escaping in JSON
+/// have been escaped.  The trailing colon is appended as well.
+final class JSONNameExtension: SchemaExtension {
+
+    static let uniqueIndex = schemaExtensionUniqueIndexCounter.getNextIndex()
+
+    let name: [UInt8]?
+
+    init(schema: Schema) throws {
+        let jsonName = schema.getTrait(JSONNameTrait.self)?.name ?? schema.id.member
+        self.name = try jsonName.map { try Serializer.writeKey(name: $0) }
+    }
+}

Sources/SmithyJSON/MemberNameExtension.swift

@@ -0,0 +1,34 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@_spi(SchemaBasedSerde)
+import class Smithy.Schema
+@_spi(SchemaBasedSerde)
+import protocol Smithy.SchemaExtension
+@_spi(SchemaBasedSerde)
+import var Smithy.schemaExtensionUniqueIndexCounter
+
+/// Stores the UTF-8 bytes for the JSON string representation of this member's JSON name.  The
+/// jsonName trait is ignored and the member name is used.
+///
+/// Bytes may be reused when this member is sent again, increasing performance by not requiring
+/// recalculation of the JSON bytes.  The tradeoff is the extra memory consumed to store the bytes
+/// in this schema extension.
+///
+/// Bytes include leading & trailing double-quotes, and all characters requiring escaping in JSON
+/// have been escaped.  The trailing colon is appended as well.
+final class MemberNameExtension: SchemaExtension {
+
+    static let uniqueIndex = schemaExtensionUniqueIndexCounter.getNextIndex()
+
+    let name: [UInt8]?
+
+    init(schema: Schema) throws {
+        let memberName = schema.id.member
+        self.name = try memberName.map { try Serializer.writeKey(name: $0) }
+    }
+}

Sources/SmithyJSON/Serializer.swift

@@ -8,8 +8,6 @@
 import struct Foundation.Data
 import struct Foundation.Date
 @_spi(SchemaBasedSerde)
-import class Smithy.JSONNameTrait
-@_spi(SchemaBasedSerde)
 import enum Smithy.Prelude
 @_spi(SchemaBasedSerde)
 import class Smithy.Schema
@@ -199,23 +197,29 @@ public final class Serializer: ShapeSerializer {
                 copyStartIndex = utf8view.index(after: index)
                 switch byte {
                 case Self.doubleQuote:
-                    appendEscaped(ascii: Self.doubleQuote)
+                    _data.append(contentsOf: [Self.backslash, Self.doubleQuote])
                 case Self.backslash:
-                    appendEscaped(ascii: Self.backslash)
+                    _data.append(contentsOf: [Self.backslash, Self.backslash])
                 case Self.backspace:
-                    appendEscaped(ascii: Self.b)
+                    _data.append(contentsOf: [Self.backslash, Self.b])
                 case Self.formFeed:
-                    appendEscaped(ascii: Self.f)
+                    _data.append(contentsOf: [Self.backslash, Self.f])
                 case Self.lineFeed:
-                    appendEscaped(ascii: Self.n)
+                    _data.append(contentsOf: [Self.backslash, Self.n])
                 case Self.cr:
-                    appendEscaped(ascii: Self.r)
+                    _data.append(contentsOf: [Self.backslash, Self.r])
                 case Self.tab:
-                    appendEscaped(ascii: Self.t)
+                    _data.append(contentsOf: [Self.backslash, Self.t])
                 case 0..<0x20:
                     // Any C0 control without a short form must be \u00XX-escaped (RFC 8259)
-                    appendEscaped(ascii: Self.u)
-                    appendHexByte(ascii: byte)
+                    _data.append(contentsOf: [
+                        Self.backslash,
+                        Self.u,
+                        Self.zero,
+                        Self.zero,
+                        Self.digits[Int(byte >> 4)],
+                        Self.digits[Int(byte & 0x0F)],
+                    ])
                 default:
                     break
                 }
@@ -299,38 +303,39 @@ public final class Serializer: ShapeSerializer {
         }
         self._needsComma = true
 
-        // If this is a member of a structure or union, write the key string and a colon.
-        // Never lead the key with a comma since it was just written above.
-        if schema.containerType == .structure || schema.containerType == .union, let key = try objectKey(for: schema) {
-            let savedNeedsComma = self._needsComma
-            self._needsComma = false
-            try writeString(Smithy.Prelude.stringSchema, key)
-            self._needsComma = savedNeedsComma
-            _data.append(Self.colon)
-        }
-    }
-
-    private func objectKey(for memberSchema: Schema) throws -> String? {
-        // Get jsonName, if present, for restJson.  Otherwise just the member name.
-        return if usesJSONNameTrait, let jsonName = memberSchema.getTrait(JSONNameTrait.self)?.name {
-            jsonName
+        // If this is a member of a structure or union, write the key, which includes a
+        // double-quoted JSON string plus a colon.
+        // The MemberNameExtension or JSONNameExtension is used to calculate then store
+        // the key for future use.
+        guard schema.containerType == .structure || schema.containerType == .union else { return }
+        let name = if usesJSONNameTrait {
+            try schema.getOrCreateExtension(JSONNameExtension.self).name
         } else {
-            memberSchema.id.member
+            try schema.getOrCreateExtension(MemberNameExtension.self).name
         }
+        guard let name else { return }
+        _data.append(contentsOf: name)
     }
 
-    private func appendEscaped(ascii: UInt8) {
-        _data.append(contentsOf: [Self.backslash, ascii])
-    }
+    static let digits: [UInt8] = "0123456789abcdef".compactMap { $0.asciiValue }
 
-    private func appendHexByte(ascii: UInt8) {
-        _data.append(contentsOf: [
-            Self.zero,
-            Self.zero,
-            Self.digits[Int(ascii >> 4)],
-            Self.digits[Int(ascii & 0x0F)],
-        ])
+    /// Encodes the passed name as the UTF-8 bytes for a JSON object key.
+    ///
+    /// The returned bytes include the enclosing double-quotes and a trailing colon, and any
+    /// characters requiring escaping in JSON have been escaped.  Used by the schema extensions
+    /// that cache a member's serialized key; see ``MemberNameExtension`` & ``JSONNameExtension``.
+    /// - Parameter name: The key name to be encoded.
+    /// - Returns: The UTF-8 bytes for the JSON representation of the key.
+    static func writeKey(name: String) throws -> [UInt8] {
+        // Create a Serializer & use it to write the key to a JSON string,
+        // surrounded by double quotes.
+        let serializer = Serializer(usesJSONNameTrait: false)
+        try serializer.writeString(Smithy.Prelude.stringSchema, name)
+
+        // Get the data, append a colon (UTF-8 58) to the end, and create a new array
+        // to trim extra capacity and flatten.
+        var data = try serializer.data
+        data.append(Self.colon)
+        return [UInt8](data)
     }
-
-    static let digits: [UInt8] = "0123456789abcdef".compactMap { $0.asciiValue }
 }