Add the `NodeReference.neighbors` implementation (#319)

The implementation is adapted from our previous implementation, but has
been refactored to be more appropriate for our generator-function
approach.

The unit tests comprehensively explore a simple example graph, testing
the following cases:
- Direction filtering (IN/OUT/ANY/unspecified)
- Node types
- Edge types
- Self-loop edges
- Dangling edges
- Removed edges don't appear

Test plan: Carefully inspect the added unit tests.

Paired with @wchargin

src/core2/graph.js

@@ -4,6 +4,7 @@ import deepEqual from "lodash.isequal";
 import type {Address} from "./address";
 import {AddressMap} from "./address";
 import type {Compatible} from "../util/compat";
+import stringify from "json-stable-stringify";
 
 export type Node<NR: NodeReference, NP: NodePayload> = {|
   +ref: NR,
@@ -160,6 +161,10 @@ export class Graph {
     if (!indexedEdge) {
       return undefined;
     }
+    return this._upgradeIndexedEdge(indexedEdge);
+  }
+
+  _upgradeIndexedEdge(indexedEdge: IndexedEdge): Edge<any> {
     return {
       address: indexedEdge.address,
       src: this._nodes[indexedEdge.srcIndex].address,
@@ -174,12 +179,9 @@ export class Graph {
    * If filter is provided, it will return only edges with the requested type.
    */
   *edges(options?: PluginFilter): Iterator<Edge<any>> {
-    let edges = this._edges.getAll().map((indexedEdge) => ({
+    let edges = this._edges
-      address: indexedEdge.address,
+      .getAll()
-      src: this._nodes[indexedEdge.srcIndex].address,
+      .map((indexedEdge) => this._upgradeIndexedEdge(indexedEdge));
-      dst: this._nodes[indexedEdge.dstIndex].address,
-      payload: indexedEdge.payload,
-    }));
     const filter = addressFilterer(options);
     for (const edge of edges) {
       if (filter(edge.address)) {
@@ -433,10 +435,51 @@ class InternalReference implements NodeReference {
     return this._graph._nodes[indexDatum.index].node;
   }
 
-  neighbors(
+  *neighbors(
     options?: NeighborsOptions
   ): Iterator<{|+ref: NodeReference, +edge: Edge<any>|}> {
-    const _ = options;
+    const indexDatum = this._graph._nodeIndices.get(this._address);
-    throw new Error("Not implemented");
+    if (indexDatum == null) {
+      return;
+    }
+    const nodeIndex = indexDatum.index;
+
+    const direction = (options && options.direction) || "ANY";
+    const edgeFilter =
+      options == null ? (_) => true : addressFilterer(options.edge);
+    const nodeFilter =
+      options == null ? (_) => true : addressFilterer(options.node);
+
+    const graph = this._graph;
+    const adjacencies = [];
+    if (direction === "ANY" || direction === "IN") {
+      adjacencies.push({list: graph._inEdges[nodeIndex], direction: "IN"});
+    }
+    if (direction === "ANY" || direction === "OUT") {
+      adjacencies.push({list: graph._outEdges[nodeIndex], direction: "OUT"});
+    }
+
+    for (const adjacency of adjacencies) {
+      for (const edgeAddress of adjacency.list) {
+        const indexedEdge = graph._edges.get(edgeAddress);
+        if (indexedEdge == null) {
+          throw new Error(
+            `Edge at address ${stringify(edgeAddress)} does not exist`
+          );
+        }
+        if (direction === "ANY" && adjacency.direction === "IN") {
+          if (indexedEdge.srcIndex === indexedEdge.dstIndex) {
+            continue;
+          }
+        }
+        const edge = graph._upgradeIndexedEdge(indexedEdge);
+        const ref = graph.ref(
+          adjacency.direction === "IN" ? edge.src : edge.dst
+        );
+        if (edgeFilter(edge.address) && nodeFilter(ref.address())) {
+          yield {edge, ref};
+        }
+      }
+    }
   }
 }

src/core2/graph.test.js

@@ -2,7 +2,7 @@
 import stringify from "json-stable-stringify";
 import sortBy from "lodash.sortby";
 
-import type {Node} from "./graph";
+import type {Node, Edge, NodeReference} from "./graph";
 import {DelegateNodeReference, Graph} from "./graph";
 
 import {
@@ -71,6 +71,138 @@ describe("graph", () => {
       expect(() => graph.ref((null: any))).toThrow("null");
       expect(() => graph.ref((undefined: any))).toThrow("undefined");
     });
+    describe("neighbors", () => {
+      // Note: The tests share more state within this block than in the rest of the code.
+      // It should be fine so long as neighbors never mutates the graph.
+      const edge = (id, src, dst, type = "EDGE") => ({
+        address: {id, owner: {plugin: EXAMPLE_PLUGIN_NAME, type}},
+        src: src.address(),
+        dst: dst.address(),
+        payload: {},
+      });
+      // A cute little diagram:
+      // A>B = edge from A to B
+      // Nodes in the graph: bar, foo, isolated
+      // (Parens) = Repeated or absent node
+      //
+      //    bar
+      //     V
+      //    foo  > (foo)
+      //     V
+      //   (absent)   isolated
+      const foo = new FooPayload();
+      const bar = new BarPayload(1, "hello");
+      const isolated = new BarPayload(666, "ghost");
+      const absent = new BarPayload(404, "hello");
+
+      const bar_foo = edge("bar_foo", bar, foo);
+      const foo_foo = edge("foo_foo", foo, foo, "SELF");
+      const foo_absent = edge("foo_absent", foo, absent);
+      const phantomEdge = edge("spooky", foo, isolated);
+      const graph = newGraph()
+        .addNode(bar)
+        .addNode(foo)
+        .addNode(isolated)
+        .addEdge(bar_foo)
+        .addEdge(foo_foo)
+        .addEdge(foo_absent)
+        .addEdge(phantomEdge)
+        .removeEdge(phantomEdge.address);
+
+      const refFor = (x) => graph.ref(x.address());
+
+      const fooNeighbor = {
+        bar: {edge: bar_foo, ref: refFor(bar)},
+        absent: {edge: foo_absent, ref: refFor(absent)},
+        foo: {edge: foo_foo, ref: refFor(foo)},
+      };
+
+      function expectNeighborsEqual(
+        actual: Iterable<{|+edge: Edge<any>, +ref: NodeReference|}>,
+        expected: {|+edge: Edge<any>, +ref: NodeReference|}[]
+      ) {
+        const sort = (xs) => sortBy(xs, (x) => stringify(x.edge.address));
+        expect(sort(Array.from(actual))).toEqual(sort(expected));
+      }
+
+      it("ref in empty graph has no neighbors", () => {
+        const ref = newGraph().ref(foo.address());
+        expectNeighborsEqual(ref.neighbors(), []);
+      });
+      it("graph with no edges has no neighbors", () => {
+        const g = newGraph()
+          .addNode(foo)
+          .addNode(bar);
+        const ref = g.ref(foo.address());
+        expectNeighborsEqual(ref.neighbors(), []);
+      });
+      it("finds neighbors for an absent node", () => {
+        expectNeighborsEqual(refFor(absent).neighbors(), [
+          {edge: foo_absent, ref: refFor(foo)},
+        ]);
+      });
+      describe("filters by direction:", () => {
+        [
+          ["IN", [fooNeighbor.bar, fooNeighbor.foo]],
+          ["OUT", [fooNeighbor.absent, fooNeighbor.foo]],
+          ["ANY", [fooNeighbor.bar, fooNeighbor.absent, fooNeighbor.foo]],
+          [
+            "unspecified",
+            [fooNeighbor.bar, fooNeighbor.absent, fooNeighbor.foo],
+          ],
+        ].forEach(([direction, expectedNeighbors]) => {
+          it(direction, () => {
+            const options = direction === "unspecified" ? {} : {direction};
+            expectNeighborsEqual(
+              refFor(foo).neighbors(options),
+              expectedNeighbors
+            );
+          });
+        });
+      });
+      describe("filters edges by type:", () => {
+        [
+          ["EDGE", [fooNeighbor.bar, fooNeighbor.absent]],
+          ["SELF", [fooNeighbor.foo]],
+          [
+            "unspecified",
+            [fooNeighbor.bar, fooNeighbor.absent, fooNeighbor.foo],
+          ],
+        ].forEach(([type, expectedNeighbors]) => {
+          it(type, () => {
+            const options =
+              type === "unspecified"
+                ? {}
+                : {edge: {plugin: EXAMPLE_PLUGIN_NAME, type}};
+            expectNeighborsEqual(
+              refFor(foo).neighbors(options),
+              expectedNeighbors
+            );
+          });
+        });
+      });
+      describe("filters nodes by type:", () => {
+        [
+          ["FOO", [fooNeighbor.foo]],
+          ["BAR", [fooNeighbor.bar, fooNeighbor.absent]],
+          [
+            "unspecified",
+            [fooNeighbor.bar, fooNeighbor.absent, fooNeighbor.foo],
+          ],
+        ].forEach(([type, expectedNeighbors]) => {
+          it(type, () => {
+            const options =
+              type === "unspecified"
+                ? {}
+                : {node: {plugin: EXAMPLE_PLUGIN_NAME, type}};
+            expectNeighborsEqual(
+              refFor(foo).neighbors(options),
+              expectedNeighbors
+            );
+          });
+        });
+      });
+    });
   });
 
   describe("node", () => {