Move logic to transformer

Author: JeremyMoeglich

src/accessDeep.ts

@@ -30,28 +30,6 @@ const getNthKey = (
   return indexed[n];
 };
 
-const rememberIndexedCollectionKeys = (
-  before: unknown,
-  after: unknown,
-  context: AccessDeepContext
-) => {
-  if (!isArray(before)) {
-    return;
-  }
-
-  if (isSet(after)) {
-    context.set(after, before.slice());
-    return;
-  }
-
-  if (isMap(after)) {
-    context.set(
-      after,
-      before.map(entry => (isArray(entry) ? entry[0] : undefined))
-    );
-  }
-};
-
 function validatePath(path: (string | number)[]) {
   if (includes(path, '__proto__')) {
     throw new Error('__proto__ is not allowed as a property');
@@ -99,15 +77,13 @@ export const getDeep = (
 export const setDeep = (
   object: any,
   path: (string | number)[],
-  mapper: (v: any) => any,
+  mapper: (v: any, context: AccessDeepContext) => any,
   context: AccessDeepContext
 ): any => {
   validatePath(path);
 
   if (path.length === 0) {
-    const mapped = mapper(object);
-    rememberIndexedCollectionKeys(object, mapped, context);
-    return mapped;
+    return mapper(object, context);
   }
 
   let parent = object;
@@ -148,14 +124,12 @@ export const setDeep = (
 
   if (isArray(parent)) {
     const oldValue = parent[+lastKey];
-    const newValue = mapper(oldValue);
+    const newValue = mapper(oldValue, context);
     parent[+lastKey] = newValue;
-    rememberIndexedCollectionKeys(oldValue, newValue, context);
   } else if (isPlainObject(parent)) {
     const oldValue = parent[lastKey];
-    const newValue = mapper(oldValue);
+    const newValue = mapper(oldValue, context);
     parent[lastKey] = newValue;
-    rememberIndexedCollectionKeys(oldValue, newValue, context);
   }
 
   if (isSet(parent)) {
@@ -166,8 +140,7 @@ export const setDeep = (
     }
     const oldValue = indexed[row];
 
-    const newValue = mapper(oldValue);
-    rememberIndexedCollectionKeys(oldValue, newValue, context);
+    const newValue = mapper(oldValue, context);
 
     if (oldValue !== newValue) {
       if (row < parent.size) {
@@ -191,8 +164,7 @@ export const setDeep = (
 
     switch (type) {
       case 'key': {
-        const newKey = mapper(keyToRow);
-        rememberIndexedCollectionKeys(keyToRow, newKey, context);
+        const newKey = mapper(keyToRow, context);
         parent.set(newKey, parent.get(keyToRow));
 
         if (!isVirtualRow && newKey !== keyToRow) {
@@ -205,8 +177,7 @@ export const setDeep = (
 
       case 'value': {
         const oldValue = parent.get(keyToRow);
-        const newValue = mapper(oldValue);
-        rememberIndexedCollectionKeys(oldValue, newValue, context);
+        const newValue = mapper(oldValue, context);
         parent.set(keyToRow, newValue);
         break;
       }

src/plainer.ts

@@ -71,7 +71,13 @@ export function applyValueAnnotations(
   traverse(
     annotations,
     (type, path) => {
-      plain = setDeep(plain, path, v => untransformValue(v, type, superJson), context);
+      plain = setDeep(
+        plain,
+        path,
+        (v, mapperContext) =>
+          untransformValue(v, type, superJson, mapperContext),
+        context
+      );
     },
     version
   );

src/transformer.ts

@@ -16,6 +16,7 @@ import {
 } from './is.js';
 import { findArr } from './util.js';
 import SuperJSON from './index.js';
+import type { AccessDeepContext } from './accessDeep.js';
 
 export type PrimitiveTypeAnnotation = 'number' | 'undefined' | 'bigint';
 
@@ -40,7 +41,11 @@ function simpleTransformation<I, O, A extends SimpleTypeAnnotation>(
   isApplicable: (v: any, superJson: SuperJSON) => v is I,
   annotation: A,
   transform: (v: I, superJson: SuperJSON) => O,
-  untransform: (v: O, superJson: SuperJSON) => I
+  untransform: (
+    v: O,
+    superJson: SuperJSON,
+    context: AccessDeepContext
+  ) => I
 ) {
   return {
     isApplicable,
@@ -127,13 +132,24 @@ const simpleRules = [
     // (sets only exist in es6+)
     // eslint-disable-next-line es5/no-es6-methods
     v => [...v.values()],
-    v => new Set(v)
+    (v, _, context) => {
+      const untransformed = new Set(v);
+      context?.set(untransformed, v.slice());
+      return untransformed;
+    }
   ),
   simpleTransformation(
     isMap,
     'map',
     v => [...v.entries()],
-    v => new Map(v)
+    (v, _, context) => {
+      const untransformed = new Map(v);
+      context?.set(
+        untransformed,
+        v.map(entry => (isArray(entry) ? entry[0] : undefined))
+      );
+      return untransformed;
+    }
   ),
 
   simpleTransformation<number, 'NaN' | 'Infinity' | '-Infinity', 'number'>(
@@ -345,7 +361,8 @@ simpleRules.forEach(rule => {
 export const untransformValue = (
   json: any,
   type: TypeAnnotation,
-  superJson: SuperJSON
+  superJson: SuperJSON,
+  context: AccessDeepContext
 ) => {
   if (isArray(type)) {
     switch (type[0]) {
@@ -366,6 +383,6 @@ export const untransformValue = (
       throw new Error('Unknown transformation: ' + type);
     }
 
-    return transformation.untransform(json as never, superJson);
+    return transformation.untransform(json as never, superJson, context);
   }
 };