diff --git a/src/app/credExplorer/pagerankTable/aggregate.js b/src/app/credExplorer/pagerankTable/aggregate.js index 1664e25..c2594e0 100644 --- a/src/app/credExplorer/pagerankTable/aggregate.js +++ b/src/app/credExplorer/pagerankTable/aggregate.js @@ -2,6 +2,7 @@ import sortBy from "lodash.sortby"; import stringify from "json-stable-stringify"; +import * as MapUtil from "../../../util/map"; import {NodeTrie, EdgeTrie} from "../../../core/trie"; import type {NodeType, EdgeType} from "../../adapters/pluginAdapter"; import type {ScoredConnection} from "../../../core/attribution/pagerankNodeDecomposition"; @@ -53,11 +54,7 @@ export function aggregateByNodeType( const nodeTypeToConnections = new Map(); for (const x of xs) { const type = typeTrie.getLast(x.source); - const connections = nodeTypeToConnections.get(type) || []; - if (connections.length === 0) { - nodeTypeToConnections.set(type, connections); - } - connections.push(x); + MapUtil.pushValue(nodeTypeToConnections, type, x); } const aggregations: NodeAggregation[] = []; for (const [ @@ -109,11 +106,7 @@ export function aggregateByConnectionType( } const edge = x.connection.adjacency.edge; const type = typeTrie.getLast(edge.address); - const connections = relevantMap.get(type) || []; - if (connections.length === 0) { - relevantMap.set(type, connections); - } - connections.push(x); + MapUtil.pushValue(relevantMap, type, x); } function createAggregation( diff --git a/src/util/map.js b/src/util/map.js index b4646d0..f0865c9 100644 --- a/src/util/map.js +++ b/src/util/map.js @@ -146,3 +146,19 @@ export function merge( } return result; } + +/** + * Given a map whose values are arrays, push an element onto the array + * corresponding to the given key. If the key is not in the map, first + * insert it with value a new empty array. + * + * If the key is already in the map, its value will be mutated, not + * replaced. + */ +export function pushValue(map: Map, key: K, value: V): void { + let arr = map.get(key); + if (arr == null) { + map.set(key, (arr = [])); + } + arr.push(value); +} diff --git a/src/util/map.test.js b/src/util/map.test.js index 577fd2d..7c7ae50 100644 --- a/src/util/map.test.js +++ b/src/util/map.test.js @@ -294,4 +294,27 @@ describe("util/map", () => { ]); }); }); + describe("pushValue", () => { + it("works when the map has no matching key", () => { + const map = new Map(); + MapUtil.pushValue(map, "foo", 3); + expect(map).toEqual(new Map().set("foo", [3])); + }); + it("works when the map has a matching key", () => { + const map = new Map().set("foo", [3]); + MapUtil.pushValue(map, "foo", 4); + expect(map).toEqual(new Map().set("foo", [3, 4])); + }); + it("works when the map already has an empty array", () => { + const map = new Map().set("foo", []); + MapUtil.pushValue(map, "foo", 1); + expect(map).toEqual(new Map().set("foo", [1])); + }); + it("preserves array identity", () => { + const arr = []; + const map = new Map().set("foo", arr); + MapUtil.pushValue(map, "foo", 1); + expect(map.get("foo")).toBe(arr); + }); + }); });