chore(core/protocols): performance improvements for shape serde traversal

Author: kuhe

packages/core/src/submodules/protocols/json/JsonShapeDeserializer.ts

@@ -18,6 +18,7 @@ import type {
 import { fromBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../ConfigurableSerdeContext";
+import { deserializingStructIterator } from "../structIterator";
 import { JsonSettings } from "./JsonCodec";
 import { jsonReviver } from "./jsonReviver";
 import { parseJsonBody } from "./parseJsonBody";
@@ -69,7 +70,11 @@ export class JsonShapeDeserializer extends SerdeContextConfig implements ShapeDe
       return out;
     } else if (ns.isStructSchema() && isObject) {
       const out = {} as any;
-      for (const [memberName, memberSchema] of ns.structIterator()) {
+      for (const [memberName, memberSchema] of deserializingStructIterator(
+        ns,
+        value,
+        this.settings.jsonName ? "jsonName" : false
+      )) {
         const fromKey = this.settings.jsonName ? memberSchema.getMergedTraits().jsonName ?? memberName : memberName;
         const deserializedValue = this._read(memberSchema, (value as any)[fromKey]);
         if (deserializedValue != null) {

packages/core/src/submodules/protocols/json/JsonShapeSerializer.ts

@@ -11,6 +11,7 @@ import type {
 import { toBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../ConfigurableSerdeContext";
+import { serializingStructIterator } from "../structIterator";
 import type { JsonSettings } from "./JsonCodec";
 import { JsonReplacer } from "./jsonReplacer";
 
@@ -80,10 +81,10 @@ export class JsonShapeSerializer extends SerdeContextConfig implements ShapeSeri
       return out;
     } else if (ns.isStructSchema() && isObject) {
       const out = {} as any;
-      for (const [memberName, memberSchema] of ns.structIterator()) {
-        const targetKey = this.settings.jsonName ? memberSchema.getMergedTraits().jsonName ?? memberName : memberName;
+      for (const [memberName, memberSchema] of serializingStructIterator(ns, value)) {
         const serializableValue = this._write(memberSchema, (value as any)[memberName], ns);
         if (serializableValue !== undefined) {
+          const targetKey = this.settings.jsonName ? memberSchema.getMergedTraits().jsonName ?? memberName : memberName;
           out[targetKey] = serializableValue;
         }
       }

packages/core/src/submodules/protocols/json/experimental/SinglePassJsonShapeSerializer.ts

@@ -11,6 +11,7 @@ import type {
 import { toBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../../ConfigurableSerdeContext";
+import { serializingStructIterator } from "../../structIterator";
 import type { JsonSettings } from "../JsonCodec";
 
 /**
@@ -72,7 +73,7 @@ export class SinglePassJsonShapeSerializer extends SerdeContextConfig implements
       }
     } else if (ns.isStructSchema()) {
       b += "{";
-      for (const [name, member] of ns.structIterator()) {
+      for (const [name, member] of serializingStructIterator(ns, value)) {
         const item = (value as any)[name];
         const targetKey = this.settings.jsonName ? member.getMergedTraits().jsonName ?? name : name;
         const serializableValue = this.writeValue(member, item);

packages/core/src/submodules/protocols/query/QueryShapeSerializer.ts

@@ -12,6 +12,7 @@ import type {
 import { toBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../ConfigurableSerdeContext";
+import { serializingStructIterator } from "../structIterator";
 import type { QuerySerializerSettings } from "./QuerySerializerSettings";
 
 /**
@@ -119,7 +120,7 @@ export class QueryShapeSerializer extends SerdeContextConfig implements ShapeSer
       }
     } else if (ns.isStructSchema()) {
       if (value && typeof value === "object") {
-        for (const [memberName, member] of ns.structIterator()) {
+        for (const [memberName, member] of serializingStructIterator(ns, value)) {
           if ((value as any)[memberName] == null && !member.isIdempotencyToken()) {
             continue;
           }

packages/core/src/submodules/protocols/structIterator.spec.ts

@@ -0,0 +1,127 @@
+import { NormalizedSchema } from "@smithy/core/schema";
+import type { StaticStructureSchema } from "@smithy/types";
+import { describe, expect, test as it, vi } from "vitest";
+
+import { deserializingStructIterator, serializingStructIterator } from "./structIterator";
+
+describe("filtered struct iteration", () => {
+  class ConstructorCountingNormalizedSchema extends (NormalizedSchema as any) {
+    public static constructorCalls = 0;
+    public constructor(...args: any[]) {
+      ConstructorCountingNormalizedSchema.constructorCalls += 1;
+      super(...args);
+    }
+  }
+
+  const schema = [
+    3,
+    "ns",
+    "Widget",
+    0,
+    ["a", "b", "c", /*d,*/ "e", "f", "g", "h", "i", "j", "k", "l"],
+    [0, [0, { jsonName: "B" }], [0, { idempotencyToken: 1 }], 0, 0, 0, 0, 0, 0, 0, 0],
+  ] satisfies StaticStructureSchema;
+
+  const ns = ConstructorCountingNormalizedSchema.of(schema);
+
+  describe("in serialization", () => {
+    it("should iterate only the keys in the source object and any idempotency tokens", () => {
+      expect(
+        [
+          ...serializingStructIterator(ns, {
+            d: "d",
+          }),
+        ].map(([k]) => k)
+      ).toEqual([
+        // a is ignored because it is not present and is not an idempotency token
+        // b is ignored because it is not present and is not an idempotency token
+        "c", // c is iterated because although it is not present, it is an idempotency token
+        // d is ignored because although it is present, it is not part of the schema
+      ]);
+      expect(
+        [
+          ...serializingStructIterator(ns, {
+            a: "a",
+            b: "b",
+            c: "c",
+            d: "d",
+          }),
+        ].map(([k]) => k)
+      ).toEqual(["a", "b", "c"]);
+    });
+  });
+
+  describe("in deserialization", () => {
+    it("should only iterate the keys that exist on the source object, accounting for jsonName", () => {
+      expect(
+        [
+          ...deserializingStructIterator(
+            ns,
+            {
+              B: "B",
+              d: "d",
+            },
+            "jsonName"
+          ),
+        ].map(([k]) => k)
+      ).toEqual([
+        // a is ignored because it is not present
+        "b", // b is iterated because its jsonName counterpart is present
+        // c is ignored because it is not present in the source object.
+        //     being an idempotencyToken doesn't mean anything in deserialization.
+        // d is ignored because although it is present, it is not part of the schema
+      ]);
+
+      expect(
+        [
+          ...deserializingStructIterator(
+            ns,
+            {
+              a: "a",
+              b: "b",
+              c: "c",
+              d: "d",
+            },
+            "jsonName"
+          ),
+        ].map(([k]) => k)
+      ).toEqual(["a", "c"]);
+      expect(
+        [
+          ...deserializingStructIterator(
+            ns,
+            {
+              a: "a",
+              b: "b",
+              c: "c",
+              d: "d",
+            },
+            false
+          ),
+        ].map(([k]) => k)
+      ).toEqual(["a", "b", "c"]);
+    });
+
+    it("halts iteration once all keys from the source object have been iterated", () => {
+      vi.spyOn(NormalizedSchema.prototype, "getMergedTraits");
+      // regular iteration iterates all schema keys
+      expect([...ns.structIterator()].map(([k]) => k)).toEqual(["a", "b", "c", "e", "f", "g", "h", "i", "j", "k", "l"]);
+      expect(NormalizedSchema.prototype.getMergedTraits).toHaveBeenCalledTimes(0);
+
+      vi.resetAllMocks();
+      expect([...deserializingStructIterator(ns, { a: "a" }, "jsonName")].map(([k]) => k)).toEqual(["a"]);
+      // only 1 call because iteration halts after 'a', since the total key count was 1.
+      expect(NormalizedSchema.prototype.getMergedTraits).toHaveBeenCalledTimes(1);
+
+      vi.resetAllMocks();
+      expect([...deserializingStructIterator(ns, { a: "a", l: "l" }, "jsonName")].map(([k]) => k)).toEqual(["a", "l"]);
+      // 11 calls because iteration continues in member order, and 'l' is the last key.
+      expect(NormalizedSchema.prototype.getMergedTraits).toHaveBeenCalledTimes(11);
+
+      vi.resetAllMocks();
+      expect([...deserializingStructIterator(ns, { a: "a", l: "l" }, false)].map(([k]) => k)).toEqual(["a", "l"]);
+      // no calls because no jsonName checking is involved.
+      expect(NormalizedSchema.prototype.getMergedTraits).toHaveBeenCalledTimes(0);
+    });
+  });
+});

packages/core/src/submodules/protocols/structIterator.ts

@@ -0,0 +1,69 @@
+import { NormalizedSchema } from "@smithy/core/schema";
+import type { StaticStructureSchema } from "@smithy/types";
+
+/**
+ * @internal
+ */
+type SourceObject = Record<string, any>;
+
+/**
+ * For serialization use only.
+ * @internal
+ *
+ * @param ns - normalized schema object.
+ * @param sourceObject - source object from serialization.
+ */
+export function* serializingStructIterator(ns: NormalizedSchema, sourceObject: SourceObject) {
+  if (ns.isUnitSchema()) {
+    return;
+  }
+  const struct = ns.getSchema() as StaticStructureSchema;
+  for (let i = 0; i < struct[4].length; ++i) {
+    const key = struct[4][i];
+    const memberNs = new (NormalizedSchema as any)([struct[5][i], 0], key);
+    if (!(key in sourceObject) && !memberNs.isIdempotencyToken()) {
+      continue;
+    }
+    yield [key, memberNs];
+  }
+}
+
+/**
+ * For deserialization use only.
+ * Yields a subset of NormalizedSchema::structIterator matched to the source object keys.
+ * This is a performance optimization to avoid creation of NormalizedSchema member
+ * objects for members that are undefined in the source data object but may be numerous
+ * in the schema/model.
+ * @internal
+ *
+ * @param ns - normalized schema object.
+ * @param sourceObject - source object from deserialization.
+ * @param nameTrait - xmlName or jsonName trait to look for.
+ */
+export function* deserializingStructIterator(
+  ns: NormalizedSchema,
+  sourceObject: SourceObject,
+  nameTrait?: "xmlName" | "jsonName" | false
+) {
+  if (ns.isUnitSchema()) {
+    return;
+  }
+  const struct = ns.getSchema() as StaticStructureSchema;
+  let keysRemaining = Object.keys(sourceObject).length;
+  for (let i = 0; i < struct[4].length; ++i) {
+    if (keysRemaining === 0) {
+      break;
+    }
+    const key = struct[4][i];
+    const memberNs = new (NormalizedSchema as any)([struct[5][i], 0], key);
+    let serializationKey = key;
+    if (nameTrait) {
+      serializationKey = memberNs.getMergedTraits()[nameTrait] ?? key;
+    }
+    if (!(serializationKey in sourceObject)) {
+      continue;
+    }
+    yield [key, memberNs];
+    keysRemaining -= 1;
+  }
+}

packages/core/src/submodules/protocols/xml/XmlShapeSerializer.ts

@@ -13,6 +13,7 @@ import type {
 import { fromBase64, toBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../ConfigurableSerdeContext";
+import { serializingStructIterator } from "../structIterator";
 import { XmlSettings } from "./XmlCodec";
 
 type XmlNamespaceAttributeValuePair = [string, string] | [undefined, undefined];
@@ -86,7 +87,7 @@ export class XmlShapeSerializer extends SerdeContextConfig implements ShapeSeria
 
     const [xmlnsAttr, xmlns] = this.getXmlnsAttribute(ns, parentXmlns);
 
-    for (const [memberName, memberSchema] of ns.structIterator()) {
+    for (const [memberName, memberSchema] of serializingStructIterator(ns, value as any)) {
       const val = (value as any)[memberName];
 
       if (val != null || memberSchema.isIdempotencyToken()) {