Repository: Qiskit/qiskit

Issue:
Refactor `EquivalenceLibrary` to internally be a graph
### What should we add?

In #7211 we refactored the search algorithm in the `BasisTranslator` transpiler pass to perform a dijkstra search using retworkx. To accomplish this the pass had to generate a retworkx graph object from the `EquivalenceLibrary` on demand and then perform the dijkstra search over that graph. However, it would be better if the equivalence library itself just internally wrapped a graph object. The data in the `EquivalenceLibrary` is inherently a graph and using a graph representation for the internal data model makes sense and will also improve the performance of the basis translator because there will be no overhead for constructing the graph object.

The output of this issue should be to change the internal `_map` attribute of the `EquivalenceLibrary` with a `retworkx.PyDiGraph` object built as the basis translator does during it's basis search: https://github.com/Qiskit/qiskit-terra/blob/8a3e760ffaf6dcd06c0a0e04f13612fdefd9ab3c/qiskit/transpiler/passes/basis/basis_translator.py#L470-L502
We should maintain the same public API that exists on the class, but just change the internal data model (and likely add helper functions to access the graph object). Then after we update the `EquivalenceLibrary` in a second pull request we can update the `BasisTranslator` to leverage the internal graph object instead of creating it on demand for each search.


Discussion:
I'd be interested in working on this.
Thanks Edwin, let me assign you.  Matthew's not at work this week, so he might be slower to respond than usual.

Resolving patch:
diff --git a/qiskit/circuit/equivalence.py b/qiskit/circuit/equivalence.py
--- a/qiskit/circuit/equivalence.py
+++ b/qiskit/circuit/equivalence.py
@@ -12,6 +12,7 @@
 
 """Gate equivalence library."""
 
+import copy
 from collections import namedtuple
 
 from rustworkx.visualization import graphviz_draw  # pylint: disable=no-name-in-module
@@ -22,10 +23,9 @@
 from .parameterexpression import ParameterExpression
 
 Key = namedtuple("Key", ["name", "num_qubits"])
-
-Entry = namedtuple("Entry", ["search_base", "equivalences"])
-
 Equivalence = namedtuple("Equivalence", ["params", "circuit"])  # Ordered to match Gate.params
+NodeData = namedtuple("NodeData", ["key", "equivs"])
+EdgeData = namedtuple("EdgeData", ["index", "num_gates", "rule", "source"])
 
 
 class EquivalenceLibrary:
@@ -37,11 +37,39 @@ def __init__(self, *, base=None):
 
         Args:
             base (Optional[EquivalenceLibrary]):  Base equivalence library to
-                will be referenced if an entry is not found in this library.
+                be referenced if an entry is not found in this library.
         """
         self._base = base
 
-        self._map = {}
+        if base is None:
+            self._graph = rx.PyDiGraph()
+            self._key_to_node_index = dict()
+            self._rule_count = 0
+        else:
+            self._graph = base._graph.copy()
+            self._key_to_node_index = copy.deepcopy(base._key_to_node_index)
+            self._rule_count = base._rule_count
+
+    @property
+    def graph(self) -> rx.PyDiGraph:
+        """Return graph representing the equivalence library data.
+
+        This property should be treated as read-only as it provides
+        a reference to the internal state of the :class:`~.EquivalenceLibrary` object.
+        If the graph returned by this property is mutated it could corrupt the
+        the contents of the object. If you need to modify the output ``PyDiGraph``
+        be sure to make a copy prior to any modification.
+
+        Returns:
+            PyDiGraph: A graph object with equivalence data in each node.
+        """
+        return self._graph
+
+    def _set_default_node(self, key):
+        """Create a new node if key not found"""
+        if key not in self._key_to_node_index:
+            self._key_to_node_index[key] = self._graph.add_node(NodeData(key=key, equivs=[]))
+        return self._key_to_node_index[key]
 
     def add_equivalence(self, gate, equivalent_circuit):
         """Add a new equivalence to the library. Future queries for the Gate
@@ -62,13 +90,25 @@ def add_equivalence(self, gate, equivalent_circuit):
         _raise_if_param_mismatch(gate.params, equivalent_circuit.parameters)
 
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         equiv = Equivalence(params=gate.params.copy(), circuit=equivalent_circuit.copy())
 
-        if key not in self._map:
-            self._map[key] = Entry(search_base=True, equivalences=[])
+        target = self._set_default_node(key)
+        self._graph[target].equivs.append(equiv)
 
-        self._map[key].equivalences.append(equiv)
+        sources = {
+            Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
+            for instruction in equivalent_circuit
+        }
+        edges = [
+            (
+                self._set_default_node(source),
+                target,
+                EdgeData(index=self._rule_count, num_gates=len(sources), rule=equiv, source=source),
+            )
+            for source in sources
+        ]
+        self._graph.add_edges_from(edges)
+        self._rule_count += 1
 
     def has_entry(self, gate):
         """Check if a library contains any decompositions for gate.
@@ -80,10 +120,9 @@ def has_entry(self, gate):
             Bool: True if gate has a known decomposition in the library.
                 False otherwise.
         """
-
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
 
-        return key in self._map or (self._base.has_entry(gate) if self._base is not None else False)
+        return key in self._key_to_node_index
 
     def set_entry(self, gate, entry):
         """Set the equivalence record for a Gate. Future queries for the Gate
@@ -98,16 +137,14 @@ def set_entry(self, gate, entry):
             entry (List['QuantumCircuit']) : A list of QuantumCircuits, each
                 equivalently implementing the given Gate.
         """
-
         for equiv in entry:
             _raise_if_shape_mismatch(gate, equiv)
             _raise_if_param_mismatch(gate.params, equiv.parameters)
 
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         equivs = [Equivalence(params=gate.params.copy(), circuit=equiv.copy()) for equiv in entry]
 
-        self._map[key] = Entry(search_base=False, equivalences=equivs)
+        self._graph[self._set_default_node(key)] = NodeData(key=key, equivs=equivs)
 
     def get_entry(self, gate):
         """Gets the set of QuantumCircuits circuits from the library which
@@ -128,13 +165,30 @@ def get_entry(self, gate):
                 ordering of the StandardEquivalenceLibrary will not generally be
                 consistent across Qiskit versions.
         """
-
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         query_params = gate.params
 
         return [_rebind_equiv(equiv, query_params) for equiv in self._get_equivalences(key)]
 
+    def keys(self):
+        """Return list of keys to key to node index map.
+
+        Returns:
+            List: Keys to the key to node index map.
+        """
+        return self._key_to_node_index.keys()
+
+    def node_index(self, key):
+        """Return node index for a given key.
+
+        Args:
+            key (Key): Key to an equivalence.
+
+        Returns:
+            Int: Index to the node in the graph for the given key.
+        """
+        return self._key_to_node_index[key]
+
     def draw(self, filename=None):
         """Draws the equivalence relations available in the library.
 
@@ -166,7 +220,7 @@ def _build_basis_graph(self):
         graph = rx.PyDiGraph()
 
         node_map = {}
-        for key in self._get_all_keys():
+        for key in self._key_to_node_index:
             name, num_qubits = key
             equivalences = self._get_equivalences(key)
 
@@ -197,23 +251,13 @@ def _build_basis_graph(self):
 
         return graph
 
-    def _get_all_keys(self):
-        base_keys = self._base._get_all_keys() if self._base is not None else set()
-
-        self_keys = set(self._map.keys())
-
-        return self_keys | {
-            base_key
-            for base_key in base_keys
-            if base_key not in self._map or self._map[base_key].search_base
-        }
-
     def _get_equivalences(self, key):
-        search_base, equivalences = self._map.get(key, (True, []))
-
-        if search_base and self._base is not None:
-            return equivalences + self._base._get_equivalences(key)
-        return equivalences
+        """Get all the equivalences for the given key"""
+        return (
+            self._graph[self._key_to_node_index[key]].equivs
+            if key in self._key_to_node_index
+            else []
+        )
 
 
 def _raise_if_param_mismatch(gate_params, circuit_parameters):
diff --git a/qiskit/transpiler/passes/basis/basis_translator.py b/qiskit/transpiler/passes/basis/basis_translator.py
--- a/qiskit/transpiler/passes/basis/basis_translator.py
+++ b/qiskit/transpiler/passes/basis/basis_translator.py
@@ -26,7 +26,7 @@
 from qiskit.circuit import Gate, ParameterVector, QuantumRegister, ControlFlowOp, QuantumCircuit
 from qiskit.dagcircuit import DAGCircuit
 from qiskit.converters import circuit_to_dag, dag_to_circuit
-from qiskit.circuit.equivalence import Key
+from qiskit.circuit.equivalence import Key, NodeData
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.exceptions import TranspilerError
 
@@ -370,17 +370,24 @@ class StopIfBasisRewritable(Exception):
 class BasisSearchVisitor(rustworkx.visit.DijkstraVisitor):  # pylint: disable=no-member
     """Handles events emitted during `rustworkx.dijkstra_search`."""
 
-    def __init__(self, graph, source_basis, target_basis, num_gates_for_rule):
+    def __init__(self, graph, source_basis, target_basis):
         self.graph = graph
         self.target_basis = set(target_basis)
         self._source_gates_remain = set(source_basis)
-        self._num_gates_remain_for_rule = dict(num_gates_for_rule)
+        self._num_gates_remain_for_rule = dict()
+        save_index = -1
+        for edata in self.graph.edges():
+            if save_index == edata.index:
+                continue
+            self._num_gates_remain_for_rule[edata.index] = edata.num_gates
+            save_index = edata.index
+
         self._basis_transforms = []
         self._predecessors = dict()
         self._opt_cost_map = dict()
 
     def discover_vertex(self, v, score):
-        gate = self.graph[v]
+        gate = self.graph[v].key
         self._source_gates_remain.discard(gate)
         self._opt_cost_map[gate] = score
         rule = self._predecessors.get(gate, None)
@@ -406,22 +413,21 @@ def examine_edge(self, edge):
         if edata is None:
             return
 
-        index = edata["index"]
-        self._num_gates_remain_for_rule[index] -= 1
+        self._num_gates_remain_for_rule[edata.index] -= 1
 
-        target = self.graph[target]
+        target = self.graph[target].key
         # if there are gates in this `rule` that we have not yet generated, we can't apply
         # this `rule`. if `target` is already in basis, it's not beneficial to use this rule.
-        if self._num_gates_remain_for_rule[index] > 0 or target in self.target_basis:
+        if self._num_gates_remain_for_rule[edata.index] > 0 or target in self.target_basis:
             raise rustworkx.visit.PruneSearch  # pylint: disable=no-member
 
     def edge_relaxed(self, edge):
         _, target, edata = edge
         if edata is not None:
-            gate = self.graph[target]
-            self._predecessors[gate] = edata["rule"]
+            gate = self.graph[target].key
+            self._predecessors[gate] = edata.rule
 
-    def edge_cost(self, edge):
+    def edge_cost(self, edge_data):
         """Returns the cost of an edge.
 
         This function computes the cost of this edge rule by summing
@@ -430,19 +436,17 @@ def edge_cost(self, edge):
         will later add it.
         """
 
-        if edge is None:
+        if edge_data is None:
             # the target of the edge is a gate in the target basis,
             # so we return a default value of 1.
             return 1
 
         cost_tot = 0
-        rule = edge["rule"]
-        for instruction in rule.circuit:
+        for instruction in edge_data.rule.circuit:
             key = Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
             cost_tot += self._opt_cost_map[key]
 
-        source = edge["source"]
-        return cost_tot - self._opt_cost_map[source]
+        return cost_tot - self._opt_cost_map[edge_data.source]
 
     @property
     def basis_transforms(self):
@@ -477,62 +481,33 @@ def _basis_search(equiv_lib, source_basis, target_basis):
     if not source_basis:
         return []
 
-    all_gates_in_lib = set()
-
-    graph = rustworkx.PyDiGraph()
-    nodes_to_indices = dict()
-    num_gates_for_rule = dict()
-
-    def lazy_setdefault(key):
-        if key not in nodes_to_indices:
-            nodes_to_indices[key] = graph.add_node(key)
-        return nodes_to_indices[key]
-
-    rcounter = 0  # running sum of the number of equivalence rules in the library.
-    for key in equiv_lib._get_all_keys():
-        target = lazy_setdefault(key)
-        all_gates_in_lib.add(key)
-        for equiv in equiv_lib._get_equivalences(key):
-            sources = {
-                Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
-                for instruction in equiv.circuit
-            }
-            all_gates_in_lib |= sources
-            edges = [
-                (
-                    lazy_setdefault(source),
-                    target,
-                    {"index": rcounter, "rule": equiv, "source": source},
-                )
-                for source in sources
-            ]
-
-            num_gates_for_rule[rcounter] = len(sources)
-            graph.add_edges_from(edges)
-            rcounter += 1
-
     # This is only neccessary since gates in target basis are currently reported by
     # their names and we need to have in addition the number of qubits they act on.
-    target_basis_keys = [
-        key
-        for gate in target_basis
-        for key in filter(lambda key, name=gate: key.name == name, all_gates_in_lib)
-    ]
+    target_basis_keys = [key for key in equiv_lib.keys() if key.name in target_basis]
+
+    graph = equiv_lib.graph
+    vis = BasisSearchVisitor(graph, source_basis, target_basis_keys)
 
-    vis = BasisSearchVisitor(graph, source_basis, target_basis_keys, num_gates_for_rule)
     # we add a dummy node and connect it with gates in the target basis.
     # we'll start the search from this dummy node.
-    dummy = graph.add_node("dummy starting node")
-    graph.add_edges_from_no_data([(dummy, nodes_to_indices[key]) for key in target_basis_keys])
-    rtn = None
-    try:
-        rustworkx.digraph_dijkstra_search(graph, [dummy], vis.edge_cost, vis)
-    except StopIfBasisRewritable:
-        rtn = vis.basis_transforms
+    dummy = graph.add_node(NodeData(key="key", equivs=[("dummy starting node", 0)]))
 
-        logger.debug("Transformation path:")
-        for gate_name, gate_num_qubits, params, equiv in rtn:
-            logger.debug("%s/%s => %s\n%s", gate_name, gate_num_qubits, params, equiv)
+    try:
+        graph.add_edges_from_no_data(
+            [(dummy, equiv_lib.node_index(key)) for key in target_basis_keys]
+        )
+        rtn = None
+        try:
+            rustworkx.digraph_dijkstra_search(graph, [dummy], vis.edge_cost, vis)
+        except StopIfBasisRewritable:
+            rtn = vis.basis_transforms
+
+            logger.debug("Transformation path:")
+            for gate_name, gate_num_qubits, params, equiv in rtn:
+                logger.debug("%s/%s => %s\n%s", gate_name, gate_num_qubits, params, equiv)
+    finally:
+        # Remove dummy node in order to return graph to original state
+        graph.remove_node(dummy)
 
     return rtn
 


Test patch:
diff --git a/test/python/circuit/test_equivalence.py b/test/python/circuit/test_equivalence.py
--- a/test/python/circuit/test_equivalence.py
+++ b/test/python/circuit/test_equivalence.py
@@ -16,6 +16,8 @@
 import unittest
 import numpy as np
 
+import rustworkx as rx
+
 from qiskit.test import QiskitTestCase
 
 from qiskit.circuit import QuantumCircuit, Parameter, Gate
@@ -24,6 +26,7 @@
 from qiskit.converters import circuit_to_instruction, circuit_to_gate
 from qiskit.circuit import EquivalenceLibrary
 from qiskit.utils import optionals
+from qiskit.circuit.equivalence import Key, Equivalence, NodeData, EdgeData
 
 from ..visualization.visualization import QiskitVisualizationTestCase, path_to_diagram_reference
 
@@ -154,6 +157,47 @@ def test_has_not_entry(self):
 
         self.assertFalse(eq_lib.has_entry(OneQubitZeroParamGate()))
 
+    def test_equivalence_graph(self):
+        """Verify valid graph created by add_equivalence"""
+
+        eq_lib = EquivalenceLibrary()
+
+        gate = OneQubitZeroParamGate()
+        first_equiv = QuantumCircuit(1)
+        first_equiv.h(0)
+        eq_lib.add_equivalence(gate, first_equiv)
+
+        equiv_copy = eq_lib._get_equivalences(Key(name="1q0p", num_qubits=1))[0].circuit
+
+        egraph = rx.PyDiGraph()
+        node_wt = NodeData(
+            key=Key(name="1q0p", num_qubits=1), equivs=[Equivalence(params=[], circuit=equiv_copy)]
+        )
+
+        egraph.add_node(node_wt)
+
+        node_wt = NodeData(key=Key(name="h", num_qubits=1), equivs=[])
+        egraph.add_node(node_wt)
+
+        edge_wt = EdgeData(
+            index=0,
+            num_gates=1,
+            rule=Equivalence(params=[], circuit=equiv_copy),
+            source=Key(name="h", num_qubits=1),
+        )
+        egraph.add_edge(0, 1, edge_wt)
+
+        for node in eq_lib.graph.nodes():
+            self.assertTrue(node in egraph.nodes())
+            for edge in eq_lib.graph.edges():
+                self.assertTrue(edge in egraph.edges())
+
+        self.assertEqual(len(eq_lib.graph.nodes()), len(egraph.nodes()))
+        self.assertEqual(len(eq_lib.graph.edges()), len(egraph.edges()))
+
+        keys = {Key(name="1q0p", num_qubits=1): 0, Key(name="h", num_qubits=1): 1}.keys()
+        self.assertEqual(keys, eq_lib.keys())
+
 
 class TestEquivalenceLibraryWithBase(QiskitTestCase):
     """Test cases for EquivalenceLibrary with base library."""
@@ -207,8 +251,9 @@ def test_add_equivalence(self):
         entry = eq_lib.get_entry(gate)
 
         self.assertEqual(len(entry), 2)
-        self.assertEqual(entry[0], second_equiv)
-        self.assertEqual(entry[1], first_equiv)
+        self.assertNotEqual(entry[0], entry[1])
+        self.assertTrue(entry[0] in [first_equiv, second_equiv])
+        self.assertTrue(entry[1] in [first_equiv, second_equiv])
 
     def test_set_entry(self):
         """Verify we find only equivalences from top when explicitly set."""