Add WeightedGraph.overrideWeights (#1606)

In a few occasions in the codebase, we need the ability to take a
WeightedGraph and apply manual user overrides to its weights (keeping
the base weights wherever non-conflicting). It's actually a fairly
simple application of Weights.merge, but since it's of general utility
I'm adding it to the WeightedGraph API.

Test plan: I've added unit tests that validate its behavior; take a
look. `yarn test` passes.

src/api/loadWeightedGraph.js

@@ -2,7 +2,6 @@
 
 import * as WeightedGraph from "../core/weightedGraph";
 import {type WeightedGraph as WeightedGraphT} from "../core/weightedGraph";
-import * as Weights from "../core/weights";
 import {type Weights as WeightsT} from "../core/weights";
 import {type NodeContraction} from "../core/graph";
 import {TaskReporter} from "../util/taskReporter";
@@ -68,9 +67,5 @@ export function _combineGraphs(
 ): WeightedGraphT {
   const merged = WeightedGraph.merge(graphs);
   const contracted = contractWeightedGraph(merged, contractions);
-  const weights = Weights.merge([contracted.weights, weightsOverrides], {
-    nodeResolver: (a, b) => b,
-    edgeResolver: (a, b) => b,
-  });
-  return {graph: contracted.graph, weights};
+  return WeightedGraph.overrideWeights(contracted, weightsOverrides);
 }

src/core/weightedGraph.js

@@ -52,3 +52,27 @@ export function merge(ws: $ReadOnlyArray<WeightedGraph>): WeightedGraph {
   const weights = Weights.merge(ws.map((w) => w.weights));
   return {graph, weights};
 }
+
+/**
+ * Create a new WeightedGraph where default weights have been overriden.
+ *
+ * This takes a base WeightedGraph along with a set of "override" weights. The
+ * new graph has the union of both the base and override weights; wherever
+ * there is a conflict, the override weights will replace the base weights.
+ * This is useful in situations where we want to let the user manually specify
+ * some weights, and ensure that the user's decisions will trump any defaults.
+ *
+ * This method does not mutuate any of the original arguments. For performance
+ * reasons, it is not a full copy; the input and output WeightedGraphs have the
+ * exact same underlying Graph, which should not be modified.
+ */
+export function overrideWeights(
+  wg: WeightedGraph,
+  overrides: WeightsT
+): WeightedGraph {
+  const weights = Weights.merge([wg.weights, overrides], {
+    nodeResolver: (a, b) => b,
+    edgeResolver: (a, b) => b,
+  });
+  return {graph: wg.graph, weights};
+}

src/core/weightedGraph.test.js

@@ -1,7 +1,7 @@
 // @flow
 
 import * as Weights from "./weights";
-import {Graph} from "./graph";
+import {Graph, NodeAddress, EdgeAddress} from "./graph";
 import * as WeightedGraph from "./weightedGraph";
 import * as GraphTest from "./graphTestUtil";
 
@@ -10,6 +10,9 @@ describe("core/weightedGraph", () => {
     expect(wg1.graph.equals(wg2.graph)).toBe(true);
     expect(wg1.weights).toEqual(wg2.weights);
   }
+  const foo = GraphTest.node("foo");
+  const bar = GraphTest.node("bar");
+  const foobar = GraphTest.edge("foobar", foo, bar);
 
   describe("empty", () => {
     it("empty produces an empty WeightedGraph", () => {
@@ -43,14 +46,12 @@ describe("core/weightedGraph", () => {
       // Not attempting to validate edge case semantics here, since this is just
       // a wrapper around Graph.merge and WeightedGraph.merge; those functions
       // are tested more thoroughly.
-      const n1 = GraphTest.node("foo");
-      const n2 = GraphTest.node("bar");
-      const g1 = new Graph().addNode(n1);
-      const g2 = new Graph().addNode(n2);
+      const g1 = new Graph().addNode(foo);
+      const g2 = new Graph().addNode(bar);
       const w1 = Weights.empty();
-      w1.nodeWeights.set(n1.address, 1);
+      w1.nodeWeights.set(foo.address, 1);
       const w2 = Weights.empty();
-      w2.nodeWeights.set(n2.address, 2);
+      w2.nodeWeights.set(bar.address, 2);
       const wg1 = {graph: g1, weights: w1};
       const wg2 = {graph: g2, weights: w2};
       const g = Graph.merge([g1, g2]);
@@ -60,4 +61,39 @@ describe("core/weightedGraph", () => {
       expectEqual(wg, wg_);
     });
   });
+
+  describe("overrideWeights", () => {
+    const example = () => {
+      const graph = new Graph().addNode(foo).addNode(bar);
+      const weights = Weights.empty();
+      weights.nodeWeights.set(NodeAddress.empty, 0);
+      weights.nodeWeights.set(foo.address, 1);
+      weights.edgeWeights.set(foobar.address, {forwards: 2, backwards: 2});
+      weights.edgeWeights.set(EdgeAddress.empty, {forwards: 3, backwards: 3});
+      return {graph, weights};
+    };
+    it("has no effect if the overrides are empty", () => {
+      const g1 = example();
+      const g2 = WeightedGraph.overrideWeights(g1, Weights.empty());
+      expectEqual(g1, g2);
+    });
+    it("takes weights from base and overrides, choosing overrides on conflicts", () => {
+      const overrides = Weights.empty();
+      overrides.nodeWeights.set(foo.address, 101);
+      overrides.nodeWeights.set(bar.address, 102);
+      overrides.edgeWeights.set(foobar.address, {
+        forwards: 103,
+        backwards: 103,
+      });
+      const expected = Weights.empty();
+      expected.nodeWeights.set(NodeAddress.empty, 0);
+      expected.nodeWeights.set(foo.address, 101);
+      expected.nodeWeights.set(bar.address, 102);
+      expected.edgeWeights.set(foobar.address, {forwards: 103, backwards: 103});
+      expected.edgeWeights.set(EdgeAddress.empty, {forwards: 3, backwards: 3});
+      const actual = WeightedGraph.overrideWeights(example(), overrides)
+        .weights;
+      expect(expected).toEqual(actual);
+    });
+  });
 });