diff --git a/lib/shared/radix-tree.js b/lib/shared/radix-tree.js new file mode 100644 --- /dev/null +++ b/lib/shared/radix-tree.js @@ -0,0 +1,162 @@ +// @flow + +import invariant from 'invariant'; + +type RadixTreeLeafNode = { + leaf: true, + part: string, + values: Array, +}; +type RadixTreeBranchNode = { + leaf: false, + part: string, + children: Array>, +}; +type RadixTreeNode = RadixTreeLeafNode | RadixTreeBranchNode; + +class RadixTree { + root: RadixTreeBranchNode = { + leaf: false, + part: '', + children: [], + }; + + add(key: string, value: V) { + let node = this.root; + let partLeft = key; + while (true) { + partLeft = partLeft.substring(node.part.length); + if (node.leaf) { + // If we recurse into a leaf, that means we have an exact match + invariant( + partLeft.length === 0, + 'RadixTree should have exact match when adding to a leaf', + ); + node.values.push(value); + return; + } + // If there is a node that is a substring, we recurse into it + // Otherwise we look for a node with the same first char + let substringNode, firstCharNode; + for (const child of node.children) { + if (!child.leaf && partLeft.startsWith(child.part)) { + substringNode = child; + break; + } else if (child.leaf && partLeft === child.part) { + substringNode = child; + break; + } + if (child.part[0] === partLeft[0] && partLeft.length > 0) { + invariant( + !firstCharNode, + 'RadixTree should not have siblings that start with the same char', + ); + firstCharNode = child; + } + } + if (substringNode) { + node = substringNode; + continue; + } + invariant(!node.leaf, "Flow doesn't realize node cannot be a leaf here"); + if (!firstCharNode) { + node.children.push({ + leaf: true, + part: partLeft, + values: [value], + }); + return; + } + + let sharedStart = ''; + const minLength = Math.min(partLeft.length, firstCharNode.part.length); + for (let i = 0; i < minLength; i++) { + if (partLeft[i] !== firstCharNode.part[i]) { + break; + } + sharedStart += partLeft[i]; + } + invariant(sharedStart, 'firstCharNode should share at least one char'); + + firstCharNode.part = firstCharNode.part.substring(sharedStart.length); + + const newBranch = { + leaf: false, + part: sharedStart, + children: [ + firstCharNode, + { + leaf: true, + part: partLeft.substring(sharedStart.length), + values: [value], + }, + ], + }; + let replaced = false; + for (let i = node.children.length - 1; i >= 0; i--) { + if (node.children[i] === firstCharNode) { + node.children[i] = newBranch; + replaced = true; + break; + } + } + invariant(replaced, 'Failed to find branch to replace in RadixTree'); + break; + } + } + + getAllMatchingPrefix(prefix: string): V[] { + const result = []; + const stack = [{ node: this.root, partLeft: prefix }]; + while (stack.length > 0) { + const { node, partLeft: prevPartLeft } = stack.pop(); + const partLeft = prevPartLeft.substring(node.part.length); + for (const child of node.children) { + const leafMatch = child.part.startsWith(partLeft); + if (!child.leaf) { + const shouldConsiderBranch = + leafMatch || partLeft.startsWith(child.part); + if (shouldConsiderBranch) { + stack.push({ node: child, partLeft }); + } + continue; + } + if (!leafMatch) { + continue; + } + for (const value of child.values) { + result.push(value); + } + } + } + return result; + } + + getAllMatchingExactly(key: string): V[] { + const result = []; + const stack = [{ node: this.root, partLeft: key }]; + while (stack.length > 0) { + const { node, partLeft: prevPartLeft } = stack.pop(); + const partLeft = prevPartLeft.substring(node.part.length); + for (const child of node.children) { + const leafMatch = child.part === partLeft; + if (!child.leaf) { + const shouldConsiderBranch = partLeft.startsWith(child.part); + if (shouldConsiderBranch) { + stack.push({ node: child, partLeft }); + } + continue; + } + if (!leafMatch) { + continue; + } + for (const value of child.values) { + result.push(value); + } + } + } + return result; + } +} + +export default RadixTree; diff --git a/lib/shared/radix-tree.test.js b/lib/shared/radix-tree.test.js new file mode 100644 --- /dev/null +++ b/lib/shared/radix-tree.test.js @@ -0,0 +1,187 @@ +// @flow + +import RadixTree from './radix-tree.js'; + +type TestValue = { + +id: number, +}; + +const radixTree = new RadixTree(); + +describe('RadixTree.add', () => { + it('should add "test", "slow", and "water"', () => { + radixTree.add('test', { id: 1 }); + radixTree.add('slow', { id: 2 }); + radixTree.add('water', { id: 3 }); + expect(radixTree.root.children.length).toBe(3); + expect(radixTree.root.children[0].part).toBe('test'); + expect(radixTree.root.children[0].leaf).toBe(true); + expect(radixTree.root.children[1].part).toBe('slow'); + expect(radixTree.root.children[1].leaf).toBe(true); + expect(radixTree.root.children[2].part).toBe('water'); + expect(radixTree.root.children[2].leaf).toBe(true); + }); + it('should group "slow" and "slower"', () => { + radixTree.add('slower', { id: 4 }); + expect(radixTree.root.children.length).toBe(3); + const slowNode = radixTree.root.children[1]; + expect(slowNode.part).toBe('slow'); + expect(slowNode.leaf).toBe(false); + if (slowNode.leaf) { + return; + } + const slowChildren = slowNode.children; + expect(slowChildren.length).toBe(2); + expect(slowChildren[0].part).toBe(''); + expect(slowChildren[0].leaf).toBe(true); + expect(slowChildren[1].part).toBe('er'); + expect(slowChildren[1].leaf).toBe(true); + }); + it('should group "water" and "wat"', () => { + radixTree.add('wat', { id: 5 }); + expect(radixTree.root.children.length).toBe(3); + const watNode = radixTree.root.children[2]; + expect(watNode.part).toBe('wat'); + expect(watNode.leaf).toBe(false); + if (watNode.leaf) { + return; + } + const watChildren = watNode.children; + expect(watChildren.length).toBe(2); + expect(watChildren[0].part).toBe('er'); + expect(watChildren[0].leaf).toBe(true); + expect(watChildren[1].part).toBe(''); + expect(watChildren[1].leaf).toBe(true); + }); + it('should group "test" and "team"', () => { + radixTree.add('team', { id: 6 }); + expect(radixTree.root.children.length).toBe(3); + const teNode = radixTree.root.children[0]; + expect(teNode.part).toBe('te'); + expect(teNode.leaf).toBe(false); + if (teNode.leaf) { + return; + } + const teChildren = teNode.children; + expect(teChildren.length).toBe(2); + expect(teChildren[0].part).toBe('st'); + expect(teChildren[0].leaf).toBe(true); + expect(teChildren[1].part).toBe('am'); + expect(teChildren[1].leaf).toBe(true); + }); + it('should group "toast" and "te" (parent of "test" and "team")', () => { + radixTree.add('toast', { id: 7 }); + expect(radixTree.root.children.length).toBe(3); + const tNode = radixTree.root.children[0]; + expect(tNode.part).toBe('t'); + expect(tNode.leaf).toBe(false); + if (tNode.leaf) { + return; + } + const tChildren = tNode.children; + expect(tChildren.length).toBe(2); + expect(tChildren[1].part).toBe('oast'); + expect(tChildren[1].leaf).toBe(true); + const teNode = tChildren[0]; + expect(teNode.part).toBe('e'); + expect(teNode.leaf).toBe(false); + if (teNode.leaf) { + return; + } + const teChildren = teNode.children; + expect(teChildren[0].part).toBe('st'); + expect(teChildren[0].leaf).toBe(true); + expect(teChildren[1].part).toBe('am'); + expect(teChildren[1].leaf).toBe(true); + }); + it('should combine "slow" and "slow" into a single leaf', () => { + radixTree.add('slow', { id: 8 }); + expect(radixTree.root.children.length).toBe(3); + const slowNode = radixTree.root.children[1]; + expect(slowNode.part).toBe('slow'); + expect(slowNode.leaf).toBe(false); + if (slowNode.leaf) { + return; + } + const slowChildren = slowNode.children; + expect(slowChildren.length).toBe(2); + const slowLeafNode = slowChildren[0]; + expect(slowLeafNode.part).toBe(''); + expect(slowLeafNode.leaf).toBe(true); + if (!slowLeafNode.leaf) { + return; + } + expect(slowLeafNode.values.length).toBe(2); + expect(slowChildren[1].part).toBe('er'); + expect(slowChildren[1].leaf).toBe(true); + }); +}); + +describe('RadixTree.getAllMatchingPrefix', () => { + it('should match "t" to "toast", "test", and "team"', () => { + const results = radixTree.getAllMatchingPrefix('t'); + expect(results.length).toBe(3); + expect(results[0].id).toBe(7); + expect(results[1].id).toBe(1); + expect(results[2].id).toBe(6); + }); + it('should match "slow" to "slow", "slow", and "slower"', () => { + const results = radixTree.getAllMatchingPrefix('slow'); + expect(results.length).toBe(3); + expect(results[0].id).toBe(2); + expect(results[1].id).toBe(8); + expect(results[2].id).toBe(4); + }); + it('should match "slowe" to "slower"', () => { + const results = radixTree.getAllMatchingPrefix('slowe'); + expect(results.length).toBe(1); + expect(results[0].id).toBe(4); + }); + it('should match "wa" to "wat" and "water"', () => { + const results = radixTree.getAllMatchingPrefix('wa'); + expect(results.length).toBe(2); + expect(results[0].id).toBe(3); + expect(results[1].id).toBe(5); + }); + it('should match "water" to "water"', () => { + const results = radixTree.getAllMatchingPrefix('water'); + expect(results.length).toBe(1); + expect(results[0].id).toBe(3); + }); + it('should match "" to all results', () => { + const results = radixTree.getAllMatchingPrefix(''); + expect(results.length).toBe(8); + }); +}); + +describe('RadixTree.getAllMatchingExactly', () => { + it('should not match "t"', () => { + const results = radixTree.getAllMatchingExactly('t'); + expect(results.length).toBe(0); + }); + it('should match "slow" to "slow" and "slow"', () => { + const results = radixTree.getAllMatchingExactly('slow'); + expect(results.length).toBe(2); + expect(results[0].id).toBe(2); + expect(results[1].id).toBe(8); + }); + it('should not match "slowe"', () => { + const results = radixTree.getAllMatchingExactly('slowe'); + expect(results.length).toBe(0); + }); + it('should match "slower" to "slower"', () => { + const results = radixTree.getAllMatchingExactly('slower'); + expect(results.length).toBe(1); + expect(results[0].id).toBe(4); + }); + it('should match "wat" to "wat"', () => { + const results = radixTree.getAllMatchingExactly('wat'); + expect(results.length).toBe(1); + expect(results[0].id).toBe(5); + }); + it('should match "water" to "water"', () => { + const results = radixTree.getAllMatchingExactly('water'); + expect(results.length).toBe(1); + expect(results[0].id).toBe(3); + }); +});