changes as requested

mtth · Sep 24, 2024 · 510d272 · 510d272
1 parent 87faea2
commit 510d272
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 49 deletions.
diff --git a/lib/types.js b/lib/types.js
@@ -198,24 +198,26 @@ class Type {
       let types = schema.map((obj) => {
         return Type.forSchema(obj, opts);
       });
+      let projectionFn;
       if (!UnionType) {
         // either automatic detection or we have a projection function
         if (typeof opts.wrapUnions === 'function') {
           // we have a projection function
           try {
-          UnionType = typeof opts.wrapUnions(types) !== 'undefined'
-            // projection function yields a function, we can use an Unwrapped type
-            ? UnwrappedUnionType
-            : WrappedUnionType;
+            projectionFn = opts.wrapUnions(types);
+            UnionType = typeof projectionFn !== 'undefined'
+              // projection function yields a function, we can use an Unwrapped type
+              ? UnwrappedUnionType
+              : WrappedUnionType;
           } catch(e) {
-            throw new Error(`Union projection function errored: ${e}`);
+            throw new Error(`Error generating projection function: ${e}`);
           }
         } else {
           UnionType = isAmbiguous(types) ? WrappedUnionType : UnwrappedUnionType;
         }
       }
       LOGICAL_TYPE = logicalType;
-      type = new UnionType(types, opts);
+      type = new UnionType(types, opts, projectionFn);
     } else { // New type definition.
       type = (function (typeName) {
         let Type = TYPES[typeName];
@@ -1241,6 +1243,44 @@ UnionType.prototype._branchConstructor = function () {
   throw new Error('unions cannot be directly wrapped');
 };
 
+
+function generateProjectionIndexer(projectionFn) {
+  return (val) => {
+    const index = projectionFn(val);
+    if (typeof index !== 'number') {
+      throw new Error(`Projected index '${index}' is not valid`);
+    }
+    return index;
+  };
+}
+
+function generateDefaultIndexer() {
+  this._dynamicBranches = null;
+  this._bucketIndices = {};
+  this.types.forEach(function (type, index) {
+    if (Type.isType(type, 'abstract', 'logical')) {
+      if (!this._dynamicBranches) {
+        this._dynamicBranches = [];
+      }
+      this._dynamicBranches.push({index, type});
+    } else {
+      let bucket = getTypeBucket(type);
+      if (this._bucketIndices[bucket] !== undefined) {
+        throw new Error(`ambiguous unwrapped union: ${j(this)}`);
+      }
+      this._bucketIndices[bucket] = index;
+    }
+  }, this);
+  return (val) => {
+    let index = this._bucketIndices[getValueBucket(val)];
+    if (this._dynamicBranches) {
+      // Slower path, we must run the value through all branches.
+      index = this._getBranchIndex(val, index);
+    }
+    return index;
+  };
+}
+
 /**
  * "Natural" union type.
  *
@@ -1261,56 +1301,33 @@ UnionType.prototype._branchConstructor = function () {
  * + `map`, `record`
  */
 class UnwrappedUnionType extends UnionType {
-  constructor (schema, opts) {
+  /**
+   *
+   * @param {*} schema
+   * @param {*} opts
+   * @param {Function|undefined} projectionFn The projection function used
+   *                                          to determine the bucket for the
+   *                                          Union. Falls back to generate
+   *                                          from `wrapUnions` parameter
+   *                                          if given.
+   */
+  constructor (schema, opts, projectionFn) {
     super(schema, opts);
 
-    this._dynamicBranches = null;
-    this._bucketIndices = {};
-
-    this.projectionFunction = (val) => this._bucketIndices[getValueBucket(val)];
-    let hasWrapUnionsFn = opts && typeof opts.wrapUnions === 'function';
-    if (hasWrapUnionsFn) {
-      const projectionFunction = opts.wrapUnions(this.types);
-      if (typeof projectionFunction === 'undefined') {
-        hasWrapUnionsFn = false;
-      } else {
-        this.projectionFunction = (val) => {
-          const index = projectionFunction(val);
-          if (typeof index !== 'number' || index >= this._bucketIndices.length) {
-            throw new Error(`Projected index ${index} is not valid`);
-          }
-          return index;
-        }
+    if (!projectionFn && opts && typeof opts.wrapUnions === 'function') {
+      try {
+        projectionFn = opts.wrapUnions(this.types);
+      } catch(e) {
+        throw new Error(`Error generating projection function: ${e}`);
       }
     }
-
-    this.types.forEach(function (type, index) {
-      if (Type.isType(type, 'abstract', 'logical')) {
-        if (!this._dynamicBranches) {
-          this._dynamicBranches = [];
-        }
-        this._dynamicBranches.push({index, type});
-      } else if (!hasWrapUnionsFn) {
-        let bucket = getTypeBucket(type);
-        if (this._bucketIndices[bucket] !== undefined) {
-          throw new Error(`ambiguous unwrapped union: ${j(this)}`);
-        }
-        this._bucketIndices[bucket] = index;
-      }
-    }, this);
+    this._getIndex = projectionFn
+      ? generateProjectionIndexer(projectionFn)
+      : generateDefaultIndexer.bind(this)(this.types);
 
     Object.freeze(this);
   }
 
-  _getIndex (val) {
-    let index = this.projectionFunction(val);
-    if (this._dynamicBranches) {
-      // Slower path, we must run the value through all branches.
-      index = this._getBranchIndex(val, index);
-    }
-    return index;
-  }
-
   _getBranchIndex (any, index) {
     let logicalBranches = this._dynamicBranches;
     for (let i = 0, l = logicalBranches.length; i < l; i++) {

diff --git a/test/test_types.js b/test/test_types.js
@@ -3549,7 +3549,7 @@ suite('types', () => {
        // Ambiguous, but we have a projection function
       const Animal = Type.forSchema(animalTypes, { wrapUnions: mockWrapUnions });
       Animal.toBuffer({ meow: '🐈' });
-      assert.equal(mockWrapUnions.calls, 2);
+      assert.equal(mockWrapUnions.calls, 1);
       assert.throws(() => Animal.toBuffer({ snap: '🐊' }), /Unknown animal/)
     });