diff --git a/lib/types.js b/lib/types.js index 95a560bf..a2d243ab 100644 --- 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 index 0d99d5f3..7d5d2380 100644 --- 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/) });