support user-defined classes and custom handlers; rework and reorganize

src/add.js

@@ -0,0 +1,10 @@
+import { handlers, defaultHandler } from './handlers.js'
+
+export let add = (...classes) => addCustom(...classes, defaultHandler)
+
+export let addCustom = (...classes) => {
+	let handler = classes.pop()
+	for (let Class of classes) {
+		handlers.set(Class.prototype, handler)
+	}
+}

src/clear.js

@@ -1,3 +1,6 @@
 import { data } from './memoize.js'
 
-export default func => data.WeakMap && data.WeakMap.delete(func.original || func)
+export let clear = func => {
+	data.delete(func)
+	data.delete(func.original)
+}

src/compareSymbols.js

@@ -1,3 +1,3 @@
 let lookup = new Map()
 let get = symbol => lookup.get(symbol) || (lookup.set(symbol, lookup.size + 1), lookup.size)
-export default symbols => symbols.sort((a, b) => get(a) - get(b))
+export let compareSymbols = (a, b) => get(a) - get(b)

src/handlers.js

@@ -0,0 +1,34 @@
+import { compareSymbols } from './compareSymbols.js'
+
+export let handlers = new Map()
+
+let OBJECT_START = {}
+let OBJECT_END = {}
+let REGEXP = {}
+let DATE = {}
+let BUFFER = {}
+
+export let defaultHandler = (obj, push, recurse) => {
+	push(OBJECT_START, Object.getPrototypeOf(obj))
+	for (let property of Object.getOwnPropertyNames(obj)
+		.sort()
+		.concat(Object.getOwnPropertySymbols(obj).sort(compareSymbols))) {
+		push(property)
+		recurse(obj[property])
+	}
+	push(OBJECT_END)
+}
+
+handlers.set(Array.prototype, defaultHandler)
+
+handlers.set(Object.prototype, defaultHandler)
+
+handlers.set(null, defaultHandler)
+
+handlers.set(RegExp.prototype, (obj, push) => push(REGEXP, obj.toString()))
+
+handlers.set(Date.prototype, (obj, push) => push(DATE, obj.getTime()))
+
+if (typeof Buffer === 'function') {
+	handlers.set(Buffer.prototype, (obj, push) => push(BUFFER, obj.toString('binary')))
+}

src/index.js

@@ -1,2 +1,3 @@
-export { default as memoize } from './memoize.js'
-export { default as clear } from './clear.js'
+export { memoize } from './memoize.js'
+export { clear } from './clear.js'
+export { add, addCustom } from './add.js'

src/keys.js

@@ -1,45 +1,23 @@
-import sortSymbols from './sortSymbols.js'
+import { handlers } from './handlers.js'
 
-let ARRAY = Symbol()
-let HOLE = Symbol()
-let OBJECT = Symbol()
-let REGEXP = Symbol()
-let DATE = Symbol()
-let BUFFER = Symbol()
-
-let array
+let array = []
+let push = array.push.bind(array)
 
 let recurse = obj => {
-	switch (typeof obj === 'object' && obj !== null && Object.getPrototypeOf(obj)) {
-		case Array.prototype:
-			array.push(ARRAY, obj.length)
-			for (let i = 0; i < obj.length; i++) i in obj ? recurse(obj[i]) : array.push(HOLE)
-			break
-		case Object.prototype: {
-			let temp = [
-				...Object.getOwnPropertyNames(obj).sort(),
-				...sortSymbols(Object.getOwnPropertySymbols(obj)),
-			]
-			array.push(OBJECT, temp.length, ...temp)
-			temp.forEach(key => recurse(obj[key]))
-			break
+	if (typeof obj === 'object' && obj !== null) {
+		let handler = handlers.get(Object.getPrototypeOf(obj))
+		if (handler) {
+			handler(obj, push, recurse)
+			return
 		}
-		case RegExp.prototype:
-			array.push(REGEXP, obj.toString())
-			break
-		case Date.prototype:
-			array.push(DATE, obj.getTime())
-			break
-		case Buffer.prototype:
-			array.push(BUFFER, obj.toString('binary'))
-			break
-		default:
-			array.push(obj)
 	}
+	push(obj)
 }
 
-export default args => {
-	array = []
-	args.forEach(recurse)
+export let keys = args => {
+	array.length = 0
+	for (let arg of args) {
+		recurse(arg)
+	}
 	return array
 }

src/memoize.js

@@ -1,17 +1,29 @@
-import keys from './keys.js'
+import { keys } from './keys.js'
 
-export let data = {}
+export let data = new WeakMap()
 
-export default func => {
+export let memoize = func => {
 	let memoized = (...args) => {
-		let here = data
-		for (let key of [func, ...keys(args)]) {
-			let M = (typeof key === 'object' && key !== null) || typeof key === 'function' ? WeakMap : Map
-			if (!here[M.name]) here[M.name] = new M()
-			if (!here[M.name].has(key)) here[M.name].set(key, {})
-			here = here[M.name].get(key)
+		let here = data.get(func)
+		if (!here) {
+			data.set(func, (here = {}))
+		}
+		for (let key of keys(args)) {
+			let M = (typeof key === 'object' && key !== null) || typeof key === 'function' ? WeakMap : Map
+			let m = here[M.name]
+			if (!m) {
+				here[M.name] = m = new M()
+			}
+			here = m.get(key)
+			if (!here) {
+				m.set(key, (here = {}))
+			}
+		}
+		if ('value' in here) {
+			return here.value
+		} else {
+			return (here.value = func(...args))
 		}
-		return 'value' in here ? here.value : (here.value = func(...args))
 	}
 	memoized.original = func
 	return memoized

test.js

@@ -1,8 +1,7 @@
 let memor = require('.')
 let assert = require('assert')
 
-let calledCount = 0
-let func = () => ++calledCount
+let func = () => ({})
 let call = memor.memoize(func)
 
 assert.equal(call([]), call([]))
@@ -21,6 +20,14 @@ assert.notEqual(call([1, , 2]), call([1, 2, ,])) // eslint-disable-line no-spars
 
 assert.notEqual(call([1, , 2]), call([1, undefined, 2])) // eslint-disable-line no-sparse-arrays
 
+{
+	let a = []
+	let b = []
+	assert.equal(call(a), call(b))
+	b.x = 0
+	assert.notEqual(call(a), call(b))
+}
+
 assert.equal(call({}), call({}))
 
 {
@@ -37,6 +44,8 @@ assert.notEqual(call({}), call({ a: undefined }))
 
 assert.notEqual(call([[0, 0], [0]]), call([[0], [0, 0]]))
 
+assert.notEqual(call([[0, 0], [0]]), call([[0, 0, [0]]]))
+
 assert.equal(call({ a: 0, b: 1 }), call({ b: 1, a: 0 }))
 
 {
@@ -90,11 +99,55 @@ assert.notEqual(call(null), call())
 
 assert.notEqual(call(true), call(true, undefined))
 
-assert.notEqual(call(), memor.memoize(() => ++calledCount)())
-
 assert.equal(call(), memor.memoize(func)())
 
-assert.notEqual(call(Object.create(null)), call(Object.create(null)))
+assert.equal(call(Object.create(null)), call(Object.create(null)))
+
+assert.notEqual(call(Object.create(null)), call({}))
+
+{
+	class Foo {}
+	memor.add(Foo)
+	let foo = new Foo()
+	let r1 = call(foo)
+	assert.notEqual(call({}), r1)
+	foo.bar = 0
+	let r2 = call(foo)
+	assert.notEqual(r1, r2)
+	foo = new Foo()
+	assert.equal(call(foo), r1)
+	foo.bar = 0
+	assert.equal(call(foo), r2)
+}
+
+{
+	class Bar {}
+	let sigil = {}
+	memor.addCustom(Bar, (obj, push, recurse) => {
+		push(sigil)
+		recurse(obj.baz)
+	})
+	let bar = new Bar()
+	let r1 = call(bar)
+	bar.quuz = 0
+	assert.equal(call(bar), r1)
+	bar.baz = 0
+	let r2 = call(bar)
+	assert.notEqual(r1, r2)
+	bar = new Bar()
+	assert.equal(call(bar), r1)
+	bar.baz = 0
+	assert.equal(call(bar), r2)
+}
+
+{
+	class Foo {}
+	class Bar {}
+	memor.add(Foo, Bar)
+	assert.equal(call(new Foo()), call(new Foo()))
+	assert.equal(call(new Bar()), call(new Bar()))
+	assert.notEqual(call(new Foo()), call(new Bar()))
+}
 
 assert.equal(call.original, func)