From a95717e6b238ebcbf20e4c9a393d26dc86581fc3 Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Fri, 12 Feb 2021 14:28:12 -0800
Subject: [PATCH 1/7] add a library containing mcpe calculation utilities,
 DependencyLists and QubitMapping data structures, and implement the swap
 update algorithm

---
 recirq/quantum_chess/mcpe_utils.py            | 194 ++++++++++++++++++
 recirq/quantum_chess/mcpe_utils_test.py       | 140 +++++++++++++
 .../quantum_chess/swap_update_transformer.py  | 188 +++++++++++++++++
 .../swap_update_transformer_test.py           | 111 ++++++++++
 4 files changed, 633 insertions(+)
 create mode 100644 recirq/quantum_chess/mcpe_utils.py
 create mode 100644 recirq/quantum_chess/mcpe_utils_test.py
 create mode 100644 recirq/quantum_chess/swap_update_transformer.py
 create mode 100644 recirq/quantum_chess/swap_update_transformer_test.py

diff --git a/recirq/quantum_chess/mcpe_utils.py b/recirq/quantum_chess/mcpe_utils.py
new file mode 100644
index 00000000..c3304f4d
--- /dev/null
+++ b/recirq/quantum_chess/mcpe_utils.py
@@ -0,0 +1,194 @@
+# Copyright 2021 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities related to the maximum consecutive positive effect (mcpe) heuristic
+cost function.
+
+These are necessary for implementing the swap-based update algorithm described
+in the paper 'A Dynamic Look-Ahead Heuristic for the Qubit Mapping Problem of
+NISQ Computers'
+(https://ieeexplore.ieee.org/abstract/document/8976109).
+"""
+from collections import defaultdict, deque
+from typing import Callable, Dict, Generator, Iterable, List, Optional, Set, Tuple
+
+import cirq
+
+
+def manhattan_dist(q1: cirq.GridQubit, q2: cirq.GridQubit) -> int:
+    """Returns the Manhattan distance between two GridQubits.
+
+    On grid devices this is the shortest path length between the two qubits.
+    """
+    return abs(q1.row - q2.row) + abs(q1.col - q2.col)
+
+
+def swap_map_fn(q1: cirq.Qid, q2: cirq.Qid) -> Callable[[cirq.Qid], cirq.Qid]:
+    """Returns a function which applies the effect of swapping two qubits."""
+    swaps = {q1: q2, q2: q1}
+    return lambda q: swaps.get(q, q)
+
+
+def effect_of_swap(swap_qubits: Tuple[cirq.GridQubit, cirq.GridQubit],
+                   gate_qubits: Tuple[cirq.GridQubit, cirq.GridQubit]) -> int:
+    """Returns the net effect of a swap on the distance between a gate's qubits.
+
+    Note that this returns >0 if the distance would decrease and <0 if it would
+    increase, which is somewhat counter-intuitive.
+
+    Args:
+      swap_qubits: the pair of qubits to swap
+      gate_qubits: the pair of qubits that the gate operates on
+    """
+    gate_after = map(swap_map_fn(*swap_qubits), gate_qubits)
+    # TODO: Using manhattan distance only works for grid devices when all qubits
+    # are fair game.
+    # This will need to be updated for
+    # https://github.com/quantumlib/ReCirq/issues/129.
+    return manhattan_dist(*gate_qubits) - manhattan_dist(*gate_after)
+
+
+class QubitMapping:
+    """Data structure representing a 1:1 map between logical and physical GridQubits.
+
+    Args:
+      initial_mapping: initial logical-to-physical qubit map.
+    """
+    def __init__(self, initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {}):
+        self.logical_to_physical = initial_mapping
+        self.physical_to_logical = {v: k for k, v in initial_mapping.items()}
+
+    def insert(self, logical_q: cirq.Qid, physical_q: cirq.GridQubit) -> None:
+        """Inserts a mapping between a logical and physical qubit."""
+        self.logical_to_physical[logical_q] = physical_q
+        self.physical_to_logical[physical_q] = logical_q
+
+    def swap_physical(self, q1: cirq.GridQubit, q2: cirq.GridQubit) -> None:
+        """Updates the mapping by swapping two physical qubits."""
+        logical_q1 = self.physical_to_logical.get(q1)
+        logical_q2 = self.physical_to_logical.get(q2)
+        self.physical_to_logical[q1], self.physical_to_logical[
+            q2] = logical_q2, logical_q1
+        self.logical_to_physical[logical_q1], self.logical_to_physical[
+            logical_q2] = q2, q1
+
+    def logical(self, qubit: cirq.GridQubit) -> cirq.Qid:
+        """Returns the logical qubit for a given physical qubit."""
+        return self.physical_to_logical.get(qubit)
+
+    def physical(self, qubit: cirq.Qid) -> cirq.GridQubit:
+        """Returns the physical qubit for a given logical qubit."""
+        return self.logical_to_physical.get(qubit)
+
+
+class DependencyLists:
+    """Data structure representing the interdependencies between qubits and
+    gates in a circuit.
+
+    The DependencyLists maps qubits to linked lists of gates that depend on that
+    qubit in execution order.
+    Additionally, the DependencyLists can compute the MCPE heuristic cost
+    function for candidate qubit swaps.
+    """
+    def __init__(self, circuit: cirq.Circuit):
+        self.dependencies = defaultdict(deque)
+        for moment in circuit:
+            for operation in moment:
+                for qubit in operation.qubits:
+                    self.dependencies[qubit].append(operation)
+
+    def peek_front(self, qubit: cirq.Qid) -> Iterable[cirq.Operation]:
+        """Returns the first gate in a qubit's dependency list."""
+        return self.dependencies[qubit][0]
+
+    def pop_front(self, qubit: cirq.Qid) -> None:
+        """Removes the first gate in a qubit's dependency list."""
+        self.dependencies[qubit].popleft()
+
+    def empty(self, qubit: cirq.Qid) -> bool:
+        """Returns true iff the qubit's dependency list is empty."""
+        return qubit not in self.dependencies or not self.dependencies[qubit]
+
+    def all_empty(self) -> bool:
+        """Returns true iff all dependency lists are empty."""
+        return all(len(dlist) == 0 for dlist in self.dependencies.values())
+
+    def active_gates(self) -> Set[cirq.Operation]:
+        """Returns the currently active gates of the circuit represented by the
+        dependency lists.
+
+        The active gates are the ones which operate on qubits that have no other
+        preceding gate operations.
+        """
+        ret = set()
+        # Recomputing the active gates from scratch is less efficient than
+        # maintaining them as the dependency lists are updated (e.g. maintaining
+        # additional state for front gates, active gates, and frozen qubits).
+        # This can be optimized later if necessary.
+        for dlist in self.dependencies.values():
+            if not dlist:
+                continue
+            gate = dlist[0]
+            if gate in ret:
+                continue
+            if all(not self.empty(q) and self.peek_front(q) == gate
+                   for q in gate.qubits):
+                ret.add(gate)
+        return ret
+
+    def _maximum_consecutive_positive_effect_impl(
+            self, swap_q1: cirq.GridQubit, swap_q2: cirq.GridQubit,
+            gates: Iterable[cirq.Operation], mapping: QubitMapping) -> int:
+        """Computes the MCPE contribution from a single qubit's dependency list.
+
+        This is where the dynamic look-ahead window is applied -- the window of
+        gates that contribute to the MCPE ends after the first gate encountered
+        which would be made worse by applying the swap (the first one with
+        effect_of_swap() < 0).
+
+        Args:
+          swap_q1: the source qubit to swap
+          swap_q2: the target qubit to swap
+          gates: the dependency list of gate operations on logical qubits
+          mapping: the mapping between logical and physical qubits for gates
+        """
+        total_cost = 0
+        for gate in gates:
+            assert len(gate.qubits) <= 2
+            if len(gate.qubits) != 2:
+                # Single-qubit gates would not be affected by the swap. We can
+                # treat the change in cost as 0 for those.
+                continue
+            swap_cost = effect_of_swap((swap_q1, swap_q2),
+                                       tuple(map(mapping.physical,
+                                                 gate.qubits)))
+            if swap_cost < 0:
+                break
+            total_cost += swap_cost
+        return total_cost
+
+    def maximum_consecutive_positive_effect(self, swap_q1: cirq.GridQubit,
+                                            swap_q2: cirq.GridQubit,
+                                            mapping: QubitMapping) -> int:
+        """Computes the MCPE heuristic cost function of applying the swap to the
+        circuit represented by this set of DependencyLists.
+
+        Args:
+          swap_q1: the source qubit to swap
+          swap_q2: the target qubit to swap
+          mapping: the mapping between logical and physical qubits for gate in the dependency lists
+        """
+        return sum(
+            self._maximum_consecutive_positive_effect_impl(
+                swap_q1, swap_q2, self.dependencies[mapping.logical(q)],
+                mapping) for q in (swap_q1, swap_q2))
diff --git a/recirq/quantum_chess/mcpe_utils_test.py b/recirq/quantum_chess/mcpe_utils_test.py
new file mode 100644
index 00000000..6053e009
--- /dev/null
+++ b/recirq/quantum_chess/mcpe_utils_test.py
@@ -0,0 +1,140 @@
+from collections import deque
+
+import pytest
+import cirq
+
+import recirq.quantum_chess.mcpe_utils as mcpe
+
+
+def test_peek():
+    x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(3))
+    g = [cirq.ISWAP(x, y), cirq.ISWAP(x, z), cirq.ISWAP(y, z)]
+    dlists = mcpe.DependencyLists(cirq.Circuit(g))
+    assert dlists.peek_front(x) == g[0]
+    assert dlists.peek_front(y) == g[0]
+    assert dlists.peek_front(z) == g[1]
+
+
+def test_pop():
+    x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(3))
+    g = [cirq.ISWAP(x, y), cirq.ISWAP(x, z), cirq.ISWAP(y, z)]
+    dlists = mcpe.DependencyLists(cirq.Circuit(g))
+
+    assert dlists.peek_front(x) == g[0]
+    dlists.pop_front(x)
+    assert dlists.peek_front(x) == g[1]
+
+
+def test_empty():
+    x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(3))
+    dlists = mcpe.DependencyLists(
+        cirq.Circuit(cirq.ISWAP(x, y), cirq.ISWAP(x, z), cirq.ISWAP(y, z)))
+
+    assert not dlists.empty(x)
+    dlists.pop_front(x)
+    assert not dlists.empty(x)
+    dlists.pop_front(x)
+    assert dlists.empty(x)
+
+    assert not dlists.all_empty()
+    dlists.pop_front(y)
+    dlists.pop_front(y)
+    dlists.pop_front(z)
+    dlists.pop_front(z)
+    assert dlists.all_empty()
+
+
+def test_active_gates():
+    w, x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(4))
+    dlists = mcpe.DependencyLists(
+        cirq.Circuit(cirq.ISWAP(x, y), cirq.ISWAP(y, z), cirq.X(w)))
+
+    assert dlists.active_gates() == {cirq.ISWAP(x, y), cirq.X(w)}
+
+
+def test_physical_mapping():
+    q = list(cirq.NamedQubit(f'q{i}') for i in range(6))
+    Q = list(cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    mapping = mcpe.QubitMapping(dict(zip(q, Q)))
+    assert list(map(mapping.physical, q)) == Q
+    assert cirq.ISWAP(q[1],
+                      q[5]).transform_qubits(mapping.physical) == cirq.ISWAP(
+                          Q[1], Q[5])
+
+
+def test_swap():
+    q = list(cirq.NamedQubit(f'q{i}') for i in range(6))
+    Q = list(cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    mapping = mcpe.QubitMapping(dict(zip(q, Q)))
+
+    mapping.swap_physical(Q[0], Q[1])
+    g = cirq.CNOT(q[0], q[2])
+    assert g.transform_qubits(mapping.physical) == cirq.CNOT(Q[1], Q[2])
+
+    mapping.swap_physical(Q[2], Q[3])
+    mapping.swap_physical(Q[1], Q[4])
+    assert g.transform_qubits(mapping.physical) == cirq.CNOT(Q[4], Q[3])
+
+
+def test_mcpe_example_8():
+    # This test is example 8 from the circuit in figure 9 of
+    # https://ieeexplore.ieee.org/abstract/document/8976109.
+    q = list(cirq.NamedQubit(f'q{i}') for i in range(6))
+    Q = list(cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    mapping = mcpe.QubitMapping(dict(zip(q, Q)))
+    dlists = mcpe.DependencyLists(
+        cirq.Circuit(cirq.CNOT(q[0], q[2]), cirq.CNOT(q[5], q[2]),
+                     cirq.CNOT(q[0], q[5]), cirq.CNOT(q[4], q[0]),
+                     cirq.CNOT(q[0], q[3]), cirq.CNOT(q[5], q[0]),
+                     cirq.CNOT(q[3], q[1])))
+
+    assert dlists.maximum_consecutive_positive_effect(Q[0], Q[1], mapping) == 4
+
+
+def test_mcpe_example_9():
+    # This test is example 9 from the circuit in figures 9 and 10 of
+    # https://ieeexplore.ieee.org/abstract/document/8976109.
+    q = list(cirq.NamedQubit(f'q{i}') for i in range(6))
+    Q = list(cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    mapping = mcpe.QubitMapping(dict(zip(q, Q)))
+    dlists = mcpe.DependencyLists(
+        cirq.Circuit(cirq.CNOT(q[0], q[2]), cirq.CNOT(q[5], q[2]),
+                     cirq.CNOT(q[0], q[5]), cirq.CNOT(q[4], q[0]),
+                     cirq.CNOT(q[0], q[3]), cirq.CNOT(q[5], q[0]),
+                     cirq.CNOT(q[3], q[1])))
+
+    # At first CNOT(q0, q2) is the active gate.
+    assert dlists.active_gates() == {cirq.CNOT(q[0], q[2])}
+    # The swaps connected to either q0 or q2 to consider are:
+    #   (Q0, Q1), (Q0, Q3), (Q1, Q2), (Q2, Q5)
+    # Of these, (Q0, Q3) and (Q2, Q5) can be discarded because they would
+    # negatively impact the active CNOT(q0, q2) gate.
+    assert mcpe.effect_of_swap((Q[0], Q[3]), (Q[0], Q[2])) < 0
+    assert mcpe.effect_of_swap((Q[2], Q[5]), (Q[0], Q[2])) < 0
+    # The remaining candidate swaps are: (Q0, Q1) and (Q1, Q2)
+    # (Q0, Q1) has a higher MCPE, so it looks better to apply that one.
+    assert dlists.maximum_consecutive_positive_effect(Q[0], Q[1], mapping) == 4
+    assert dlists.maximum_consecutive_positive_effect(Q[1], Q[2], mapping) == 1
+    mapping.swap_physical(Q[0], Q[1])
+
+    # The swap-update algorithm would now advance beyond the front-most gates that
+    # now satisfy adjacency constraints after the swap -- the CNOT(q0, q2) and
+    # CNOT(q5, q2)
+    assert dlists.active_gates() == {cirq.CNOT(q[0], q[2])}
+    dlists.pop_front(q[0])
+    dlists.pop_front(q[2])
+    assert dlists.active_gates() == {cirq.CNOT(q[5], q[2])}
+    dlists.pop_front(q[5])
+    dlists.pop_front(q[2])
+
+    # Now the active gate is g2 (which is CNOT(q0, q5))
+    assert dlists.active_gates() == {cirq.CNOT(q[0], q[5])}
+    # For this active gate, the swaps to consider are:
+    #   (Q0, Q1), (Q1, Q2), (Q1, Q4), (Q2, Q5), (Q4, Q5)
+    # (Q0, Q1) can be discarded because it negatively impacts the active gate.
+    assert mcpe.effect_of_swap((Q[0], Q[1]), (Q[1], Q[5])) < 0
+    # Of the remaining candidate swaps, (Q0, Q4) has the highest MCPE.
+    assert dlists.maximum_consecutive_positive_effect(Q[1], Q[2], mapping) == 1
+    assert dlists.maximum_consecutive_positive_effect(Q[1], Q[4], mapping) == 3
+    assert dlists.maximum_consecutive_positive_effect(Q[2], Q[5], mapping) == 2
+    assert dlists.maximum_consecutive_positive_effect(Q[4], Q[5], mapping) == 2
diff --git a/recirq/quantum_chess/swap_update_transformer.py b/recirq/quantum_chess/swap_update_transformer.py
new file mode 100644
index 00000000..3ac4bdce
--- /dev/null
+++ b/recirq/quantum_chess/swap_update_transformer.py
@@ -0,0 +1,188 @@
+# Copyright 2021 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of the swap update algorithm described in the paper 'A Dynamic
+Look-Ahead Heuristic for the Qubit Mapping Problem of NISQ Computers'
+(https://ieeexplore.ieee.org/abstract/document/8976109).
+
+This transforms circuits by adding additional SWAP gates to ensure that all operations are on adjacent qubits.
+"""
+from collections import defaultdict
+from typing import Callable, Dict, Generator, Iterable, List, Optional, Tuple
+
+import cirq
+
+import recirq.quantum_chess.circuit_transformer as ct
+import recirq.quantum_chess.mcpe_utils as mcpe
+
+ADJACENCY = [(0, 1), (0, -1), (1, 0), (-1, 0)]
+
+
+def _satisfies_adjacency(gate: cirq.Operation) -> bool:
+    """Returns true iff the input gate operates on adjacent qubits.
+
+    Requires the input to be either a binary or unary operation on GridQubits.
+    """
+    assert len(gate.qubits) <= 2
+    if len(gate.qubits) < 2:
+        return True
+    q1, q2 = gate.qubits
+    return q1.is_adjacent(q2)
+
+
+def generate_swap(q1: cirq.Qid,
+                  q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
+    """Generates a single SWAP gate."""
+    yield cirq.SWAP(q1, q2)
+
+
+def generate_iswap(q1: cirq.Qid,
+                   q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
+    """Generates an ISWAP as two sqrt-ISWAP gates.
+
+    This is a suitable replacement for SWAPs (up to a phase shift) but is a more
+    efficient representation on hardware supporting sqrt-ISWAP natively (ex
+    Google Sycamore chips).
+    """
+    yield cirq.ISWAP(q1, q2)**0.5
+    yield cirq.ISWAP(q1, q2)**0.5
+
+
+class SwapUpdater:
+    """SwapUpdater runs the swap update algorithm in order to incrementally update a circuit with SWAPs.
+
+    The SwapUpdater's internal state is modified as the algorithm runs, so each instance is one-time use.
+
+    Args:
+      circuit: the circuit to be updated with additional SWAPs
+      device_qubits: the allowed set of qubits on the device. If None, behaves
+        as though operating on an unconstrained infinite device.
+      initial_mapping: the initial logical-to-physical qubit mapping, which
+        must contain an entry for every qubit in the circuit
+      swap_factory: the factory used to produce operations representing a swap
+        of two qubits
+    """
+    def __init__(self,
+                 circuit: cirq.Circuit,
+                 device_qubits: Optional[Iterable[cirq.GridQubit]],
+                 initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {},
+                 swap_factory: Callable[[cirq.Qid, cirq.Qid],
+                                        List[cirq.Operation]] = generate_swap):
+        self.device_qubits = device_qubits
+        self.dlists = mcpe.DependencyLists(circuit)
+        self.mapping = mcpe.QubitMapping(initial_mapping)
+        self.swap_factory = swap_factory
+
+    def _adjacent_qubits(
+            self,
+            qubit: cirq.GridQubit) -> Generator[cirq.GridQubit, None, None]:
+        """Generates the qubits adjacent to a given GridQubit on the device."""
+        for diff in ADJACENCY:
+            other = qubit + diff
+            if self.device_qubits is None or other in self.device_qubits:
+                yield other
+
+    def generate_candidate_swaps(
+        self, gates: Iterable[cirq.Operation]
+    ) -> Generator[Tuple[cirq.GridQubit, cirq.GridQubit], None, None]:
+        """Generates the candidate SWAPs that would have a positive effect on at
+        least one of the given physical gates.
+
+        Args:
+          gates: the list of gates to consider which operate on GridQubits
+        """
+        for gate in gates:
+            for gate_q in gate.qubits:
+                yield from (
+                    (gate_q, swap_q)
+                    for swap_q in self._adjacent_qubits(gate_q)
+                    if mcpe.effect_of_swap((gate_q, swap_q), gate.qubits) > 0)
+
+    def update_iteration(self) -> Generator[cirq.Operation, None, None]:
+        """Runs one iteration of the swap update algorithm and updates internal
+        state about the original circuit.
+
+        Returns:
+          the operations on GridQubits in the final updated circuit generated by
+          this iteration
+        """
+        # Handle the already-satisfied active gates.
+        # Those can be immediately added into the final circuit.
+        active_physical_gates = []
+        for gate in self.dlists.active_gates():
+            physical_gate = gate.transform_qubits(self.mapping.physical)
+            if _satisfies_adjacency(physical_gate):
+                # physical_gate is ready to be popped off the dependecy lists
+                # and added to the final circuit.
+                for q in gate.qubits:
+                    self.dlists.pop_front(q)
+                yield physical_gate
+            else:
+                # physical_gate needs to be fixed up with some swaps.
+                active_physical_gates.append(physical_gate)
+
+        # If all the active gates in this pass were added to the final circuit,
+        # then we have nothing left to do until we make another pass and get the
+        # newly active gates.
+        if not active_physical_gates:
+            return
+
+        candidates = set(self.generate_candidate_swaps(active_physical_gates))
+        # This could happen on irregular topologies or when some device
+        # qubits are disallowed.
+        assert candidates, 'no swaps found that will improve the circuit'
+        chosen_swap = max(
+            candidates,
+            key=lambda swap: self.dlists.maximum_consecutive_positive_effect(
+                *swap, self.mapping))
+        self.mapping.swap_physical(*chosen_swap)
+        yield from self.swap_factory(*chosen_swap)
+
+    def add_swaps(self) -> Generator[cirq.Operation, None, None]:
+        """Iterates the swap update algorithm to completion.
+
+        If the updater already completed, does nothing.
+
+        Returns:
+          the generated operations on physical GridQubits in the final circuit
+        """
+        while not self.dlists.all_empty():
+            yield from self.update_iteration()
+
+
+class SwapUpdateTransformer(ct.CircuitTransformer):
+    """Transformer that runs the SwapUpdater to transform circuits.
+
+    Args:
+      initial_mapping: the initial mapping between the circuit's logical qubits
+        and physical GridQubits. This must contain all qubits in the circuit to
+        be transformed.
+    """
+    def __init__(self, initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {}):
+        super().__init__()
+        self.initial_mapping = initial_mapping
+
+    def transform(self, circuit: cirq.Circuit) -> cirq.Circuit:
+        """Adds ISWAPs to satisfy adjacency constraints.
+
+        Args:
+          circuit: the input circuit of GridQubits to modify
+        Returns:
+          the modified circuit with ISWAPs inserted and operations mapped to
+          physical GridQubits.
+        """
+        for q in circuit.all_qubits():
+            assert q in initial_mapping
+        updater = SwapUpdater(circuit, circuit.device.qubit_set(),
+                              initial_mapping, generate_iswap)
+        return cirq.Circuit(updater.add_swaps(), device=circuit.device)
diff --git a/recirq/quantum_chess/swap_update_transformer_test.py b/recirq/quantum_chess/swap_update_transformer_test.py
new file mode 100644
index 00000000..246ca1b8
--- /dev/null
+++ b/recirq/quantum_chess/swap_update_transformer_test.py
@@ -0,0 +1,111 @@
+import pytest
+import cirq
+
+import recirq.quantum_chess.mcpe_utils as mcpe
+from recirq.quantum_chess.swap_update_transformer import SwapUpdater, SwapUpdateTransformer, generate_iswap
+import recirq.quantum_chess.quantum_moves as qm
+
+# Logical qubits q0 - q5.
+q = list(cirq.NamedQubit(f'q{i}') for i in range(6))
+
+# Qubits corresponding to figure 9 in
+# https://ieeexplore.ieee.org/abstract/document/8976109.
+# Coupling graph looks something like:
+#   Q0 = Q1 = Q2
+#   ||   ||   ||
+#   Q3 = Q4 = Q5
+FIGURE_9A_PHYSICAL_QUBITS = list(
+    cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+# Circuit from figure 9a in
+# https://ieeexplore.ieee.org/abstract/document/8976109.
+FIGURE_9A_CIRCUIT = cirq.Circuit(cirq.CNOT(q[0], q[2]), cirq.CNOT(q[5], q[2]),
+                                 cirq.CNOT(q[0], q[5]), cirq.CNOT(q[4], q[0]),
+                                 cirq.CNOT(q[0], q[3]), cirq.CNOT(q[5], q[0]),
+                                 cirq.CNOT(q[3], q[1]))
+
+
+def test_example_9_candidate_swaps():
+    Q = FIGURE_9A_PHYSICAL_QUBITS
+    initial_mapping = dict(zip(q, Q))
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping)
+
+    gate = cirq.CNOT(Q[0], Q[2])
+    candidates = list(updater.generate_candidate_swaps([gate]))
+    # The swaps connected to either q0 or q2 are:
+    #   (Q0, Q1), (Q0, Q3), (Q2, Q1), (Q2, Q5)
+    # But swapping (Q0, Q3) or (Q2, Q5) would negatively impact the distance
+    # between q0 and q2, so those swaps are discarded.
+    assert candidates == [(Q[0], Q[1]), (Q[2], Q[1])]
+
+
+def test_example_9_iterations():
+    Q = FIGURE_9A_PHYSICAL_QUBITS
+    initial_mapping = dict(zip(q, Q))
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping)
+
+    # First iteration adds a swap between Q0 and Q1.
+    assert list(updater.update_iteration()) == [cirq.SWAP(Q[0], Q[1])]
+    # Next two iterations add the active gates as-is.
+    assert list(updater.update_iteration()) == [cirq.CNOT(Q[1], Q[2])]
+    assert list(updater.update_iteration()) == [cirq.CNOT(Q[5], Q[2])]
+    # Next iteration adds a swap between Q1 and Q4.
+    assert list(updater.update_iteration()) == [cirq.SWAP(Q[1], Q[4])]
+    # Remaining gates are added as-is.
+    assert list(updater.update_iteration()) == [cirq.CNOT(Q[4], Q[5])]
+    assert list(updater.update_iteration()) == [cirq.CNOT(Q[1], Q[4])]
+    assert list(updater.update_iteration()) == [cirq.CNOT(Q[4], Q[3])]
+    # The final two gates are added in the same iteration, since they operate on
+    # mutually exclusive qubits and are both simultaneously active.
+    assert set(updater.update_iteration()) == {
+        cirq.CNOT(Q[5], Q[4]), cirq.CNOT(Q[3], Q[0])
+    }
+
+
+def test_example_9():
+    Q = FIGURE_9A_PHYSICAL_QUBITS
+    initial_mapping = dict(zip(q, Q))
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping)
+    updated_circuit = cirq.Circuit(
+        SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping).add_swaps())
+    assert updated_circuit == cirq.Circuit(cirq.SWAP(Q[0], Q[1]),
+                                           cirq.CNOT(Q[1], Q[2]),
+                                           cirq.CNOT(Q[5], Q[2]),
+                                           cirq.SWAP(Q[1], Q[4]),
+                                           cirq.CNOT(Q[4], Q[5]),
+                                           cirq.CNOT(Q[1], Q[4]),
+                                           cirq.CNOT(Q[4], Q[3]),
+                                           cirq.CNOT(Q[5], Q[4]),
+                                           cirq.CNOT(Q[3], Q[0]))
+
+
+def test_pentagonal_split_and_merge():
+    grid_3x2 = list(
+        cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    logical_qubits = list(
+        cirq.NamedQubit(f'{x}{i}') for x in ('a', 'b') for i in range(3))
+    initial_mapping = dict(zip(logical_qubits, grid_3x2))
+    a1, a2, a3, b1, b2, b3 = logical_qubits
+    circuit = cirq.Circuit(qm.normal_move(a1, b1), qm.normal_move(a2, a3),
+                           qm.merge_move(a3, b1, b3))
+
+    updated_circuit = cirq.Circuit(
+        SwapUpdater(circuit, grid_3x2, initial_mapping).add_swaps())
+    for op in updated_circuit.all_operations():
+        assert len(op.qubits) == 2
+        q1, q2 = op.qubits
+        assert q1 in grid_3x2
+        assert q2 in grid_3x2
+        assert q1.is_adjacent(q2)
+
+
+def test_with_iswaps():
+    Q = FIGURE_9A_PHYSICAL_QUBITS
+    initial_mapping = dict(zip(q, Q))
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
+                          generate_iswap)
+    # First iteration adds a swap between Q0 and Q1.
+    # generate_iswap implements that swap operation as two sqrt-iswaps.
+    assert list(updater.update_iteration()) == [
+        cirq.ISWAP(Q[0], Q[1])**0.5,
+        cirq.ISWAP(Q[0], Q[1])**0.5,
+    ]

From 6190fadeb5c2f76f46ab7c86bf98a23e000048f2 Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Tue, 16 Feb 2021 10:25:10 -0800
Subject: [PATCH 2/7] remove iswap factory and replace swap factory with
 sqrt-iswap decomposition

---
 .../quantum_chess/swap_update_transformer.py  | 26 +++++++------------
 .../swap_update_transformer_test.py           | 20 +++++++++-----
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/recirq/quantum_chess/swap_update_transformer.py b/recirq/quantum_chess/swap_update_transformer.py
index 3ac4bdce..5978c9ca 100644
--- a/recirq/quantum_chess/swap_update_transformer.py
+++ b/recirq/quantum_chess/swap_update_transformer.py
@@ -40,20 +40,11 @@ def _satisfies_adjacency(gate: cirq.Operation) -> bool:
     return q1.is_adjacent(q2)
 
 
-def generate_swap(q1: cirq.Qid,
-                  q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
-    """Generates a single SWAP gate."""
-    yield cirq.SWAP(q1, q2)
-
-
-def generate_iswap(q1: cirq.Qid,
-                   q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
-    """Generates an ISWAP as two sqrt-ISWAP gates.
-
-    This is a suitable replacement for SWAPs (up to a phase shift) but is a more
-    efficient representation on hardware supporting sqrt-ISWAP natively (ex
-    Google Sycamore chips).
-    """
+def generate_decomposed_swap(
+        q1: cirq.Qid, q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
+    """Generates a SWAP operation using sqrt-iswap gates."""
+    yield cirq.ISWAP(q1, q2)**0.5
+    yield cirq.ISWAP(q1, q2)**0.5
     yield cirq.ISWAP(q1, q2)**0.5
     yield cirq.ISWAP(q1, q2)**0.5
 
@@ -76,8 +67,9 @@ def __init__(self,
                  circuit: cirq.Circuit,
                  device_qubits: Optional[Iterable[cirq.GridQubit]],
                  initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {},
-                 swap_factory: Callable[[cirq.Qid, cirq.Qid],
-                                        List[cirq.Operation]] = generate_swap):
+                 swap_factory: Callable[
+                     [cirq.Qid, cirq.Qid],
+                     List[cirq.Operation]] = generate_decomposed_swap):
         self.device_qubits = device_qubits
         self.dlists = mcpe.DependencyLists(circuit)
         self.mapping = mcpe.QubitMapping(initial_mapping)
@@ -184,5 +176,5 @@ def transform(self, circuit: cirq.Circuit) -> cirq.Circuit:
         for q in circuit.all_qubits():
             assert q in initial_mapping
         updater = SwapUpdater(circuit, circuit.device.qubit_set(),
-                              initial_mapping, generate_iswap)
+                              initial_mapping)
         return cirq.Circuit(updater.add_swaps(), device=circuit.device)
diff --git a/recirq/quantum_chess/swap_update_transformer_test.py b/recirq/quantum_chess/swap_update_transformer_test.py
index 246ca1b8..ce66e983 100644
--- a/recirq/quantum_chess/swap_update_transformer_test.py
+++ b/recirq/quantum_chess/swap_update_transformer_test.py
@@ -2,7 +2,7 @@
 import cirq
 
 import recirq.quantum_chess.mcpe_utils as mcpe
-from recirq.quantum_chess.swap_update_transformer import SwapUpdater, SwapUpdateTransformer, generate_iswap
+from recirq.quantum_chess.swap_update_transformer import SwapUpdater, SwapUpdateTransformer, generate_decomposed_swap
 import recirq.quantum_chess.quantum_moves as qm
 
 # Logical qubits q0 - q5.
@@ -41,7 +41,8 @@ def test_example_9_candidate_swaps():
 def test_example_9_iterations():
     Q = FIGURE_9A_PHYSICAL_QUBITS
     initial_mapping = dict(zip(q, Q))
-    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping)
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
+                          lambda q1, q2: [cirq.SWAP(q1, q2)])
 
     # First iteration adds a swap between Q0 and Q1.
     assert list(updater.update_iteration()) == [cirq.SWAP(Q[0], Q[1])]
@@ -64,9 +65,11 @@ def test_example_9_iterations():
 def test_example_9():
     Q = FIGURE_9A_PHYSICAL_QUBITS
     initial_mapping = dict(zip(q, Q))
-    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping)
+    updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
+                          lambda q1, q2: [cirq.SWAP(q1, q2)])
     updated_circuit = cirq.Circuit(
-        SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping).add_swaps())
+        SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
+                    lambda q1, q2: [cirq.SWAP(q1, q2)]).add_swaps())
     assert updated_circuit == cirq.Circuit(cirq.SWAP(Q[0], Q[1]),
                                            cirq.CNOT(Q[1], Q[2]),
                                            cirq.CNOT(Q[5], Q[2]),
@@ -98,14 +101,17 @@ def test_pentagonal_split_and_merge():
         assert q1.is_adjacent(q2)
 
 
-def test_with_iswaps():
+def test_decomposed_swaps():
     Q = FIGURE_9A_PHYSICAL_QUBITS
     initial_mapping = dict(zip(q, Q))
     updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
-                          generate_iswap)
+                          generate_decomposed_swap)
     # First iteration adds a swap between Q0 and Q1.
-    # generate_iswap implements that swap operation as two sqrt-iswaps.
+    # generate_decomposed_swap implements that swap operation as four
+    # sqrt-iswaps.
     assert list(updater.update_iteration()) == [
         cirq.ISWAP(Q[0], Q[1])**0.5,
         cirq.ISWAP(Q[0], Q[1])**0.5,
+        cirq.ISWAP(Q[0], Q[1])**0.5,
+        cirq.ISWAP(Q[0], Q[1])**0.5,
     ]

From 396e94670dac408a745a749b9404444e2abb6b40 Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Tue, 16 Feb 2021 10:42:05 -0800
Subject: [PATCH 3/7] fix typo

---
 recirq/quantum_chess/swap_update_transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/recirq/quantum_chess/swap_update_transformer.py b/recirq/quantum_chess/swap_update_transformer.py
index 5978c9ca..481504d1 100644
--- a/recirq/quantum_chess/swap_update_transformer.py
+++ b/recirq/quantum_chess/swap_update_transformer.py
@@ -174,7 +174,7 @@ def transform(self, circuit: cirq.Circuit) -> cirq.Circuit:
           physical GridQubits.
         """
         for q in circuit.all_qubits():
-            assert q in initial_mapping
+            assert q in self.initial_mapping
         updater = SwapUpdater(circuit, circuit.device.qubit_set(),
-                              initial_mapping)
+                              self.initial_mapping)
         return cirq.Circuit(updater.add_swaps(), device=circuit.device)

From b350914cb7e6ce6c4cca922b09aefd6b95038c52 Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Wed, 17 Feb 2021 13:14:58 -0800
Subject: [PATCH 4/7] replace _adjacent_qubits() with GridQubit.neighbors() and
 remove SwapUpdateTransformer

---
 ..._update_transformer.py => swap_updater.py} | 42 +------------------
 ...ansformer_test.py => swap_updater_test.py} |  2 +-
 2 files changed, 2 insertions(+), 42 deletions(-)
 rename recirq/quantum_chess/{swap_update_transformer.py => swap_updater.py} (78%)
 rename recirq/quantum_chess/{swap_update_transformer_test.py => swap_updater_test.py} (97%)

diff --git a/recirq/quantum_chess/swap_update_transformer.py b/recirq/quantum_chess/swap_updater.py
similarity index 78%
rename from recirq/quantum_chess/swap_update_transformer.py
rename to recirq/quantum_chess/swap_updater.py
index 481504d1..25bc50c5 100644
--- a/recirq/quantum_chess/swap_update_transformer.py
+++ b/recirq/quantum_chess/swap_updater.py
@@ -22,11 +22,8 @@
 
 import cirq
 
-import recirq.quantum_chess.circuit_transformer as ct
 import recirq.quantum_chess.mcpe_utils as mcpe
 
-ADJACENCY = [(0, 1), (0, -1), (1, 0), (-1, 0)]
-
 
 def _satisfies_adjacency(gate: cirq.Operation) -> bool:
     """Returns true iff the input gate operates on adjacent qubits.
@@ -75,15 +72,6 @@ def __init__(self,
         self.mapping = mcpe.QubitMapping(initial_mapping)
         self.swap_factory = swap_factory
 
-    def _adjacent_qubits(
-            self,
-            qubit: cirq.GridQubit) -> Generator[cirq.GridQubit, None, None]:
-        """Generates the qubits adjacent to a given GridQubit on the device."""
-        for diff in ADJACENCY:
-            other = qubit + diff
-            if self.device_qubits is None or other in self.device_qubits:
-                yield other
-
     def generate_candidate_swaps(
         self, gates: Iterable[cirq.Operation]
     ) -> Generator[Tuple[cirq.GridQubit, cirq.GridQubit], None, None]:
@@ -97,7 +85,7 @@ def generate_candidate_swaps(
             for gate_q in gate.qubits:
                 yield from (
                     (gate_q, swap_q)
-                    for swap_q in self._adjacent_qubits(gate_q)
+                    for swap_q in gate_q.neighbors(self.device_qubits)
                     if mcpe.effect_of_swap((gate_q, swap_q), gate.qubits) > 0)
 
     def update_iteration(self) -> Generator[cirq.Operation, None, None]:
@@ -150,31 +138,3 @@ def add_swaps(self) -> Generator[cirq.Operation, None, None]:
         """
         while not self.dlists.all_empty():
             yield from self.update_iteration()
-
-
-class SwapUpdateTransformer(ct.CircuitTransformer):
-    """Transformer that runs the SwapUpdater to transform circuits.
-
-    Args:
-      initial_mapping: the initial mapping between the circuit's logical qubits
-        and physical GridQubits. This must contain all qubits in the circuit to
-        be transformed.
-    """
-    def __init__(self, initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {}):
-        super().__init__()
-        self.initial_mapping = initial_mapping
-
-    def transform(self, circuit: cirq.Circuit) -> cirq.Circuit:
-        """Adds ISWAPs to satisfy adjacency constraints.
-
-        Args:
-          circuit: the input circuit of GridQubits to modify
-        Returns:
-          the modified circuit with ISWAPs inserted and operations mapped to
-          physical GridQubits.
-        """
-        for q in circuit.all_qubits():
-            assert q in self.initial_mapping
-        updater = SwapUpdater(circuit, circuit.device.qubit_set(),
-                              self.initial_mapping)
-        return cirq.Circuit(updater.add_swaps(), device=circuit.device)
diff --git a/recirq/quantum_chess/swap_update_transformer_test.py b/recirq/quantum_chess/swap_updater_test.py
similarity index 97%
rename from recirq/quantum_chess/swap_update_transformer_test.py
rename to recirq/quantum_chess/swap_updater_test.py
index ce66e983..8e0177e8 100644
--- a/recirq/quantum_chess/swap_update_transformer_test.py
+++ b/recirq/quantum_chess/swap_updater_test.py
@@ -2,7 +2,7 @@
 import cirq
 
 import recirq.quantum_chess.mcpe_utils as mcpe
-from recirq.quantum_chess.swap_update_transformer import SwapUpdater, SwapUpdateTransformer, generate_decomposed_swap
+from recirq.quantum_chess.swap_updater import SwapUpdater, generate_decomposed_swap
 import recirq.quantum_chess.quantum_moves as qm
 
 # Logical qubits q0 - q5.

From f1a38827277c2d1f44a93853f60e7d927d8a3b98 Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Fri, 19 Feb 2021 13:05:38 -0800
Subject: [PATCH 5/7] replace assert's with ValueError's, update manhattan
 distance TODO, and add test coverage for previously untested public functions

---
 recirq/quantum_chess/mcpe_utils.py        | 19 +++++------
 recirq/quantum_chess/mcpe_utils_test.py   | 37 ++++++++++++++++++++
 recirq/quantum_chess/swap_updater.py      | 25 ++++++++------
 recirq/quantum_chess/swap_updater_test.py | 41 +++++++++++++++++------
 4 files changed, 90 insertions(+), 32 deletions(-)

diff --git a/recirq/quantum_chess/mcpe_utils.py b/recirq/quantum_chess/mcpe_utils.py
index c3304f4d..70cda477 100644
--- a/recirq/quantum_chess/mcpe_utils.py
+++ b/recirq/quantum_chess/mcpe_utils.py
@@ -51,10 +51,11 @@ def effect_of_swap(swap_qubits: Tuple[cirq.GridQubit, cirq.GridQubit],
       gate_qubits: the pair of qubits that the gate operates on
     """
     gate_after = map(swap_map_fn(*swap_qubits), gate_qubits)
-    # TODO: Using manhattan distance only works for grid devices when all qubits
-    # are fair game.
-    # This will need to be updated for
-    # https://github.com/quantumlib/ReCirq/issues/129.
+    # TODO(https://github.com/quantumlib/ReCirq/issues/149):
+    # Using manhattan distance only works for grid devices when all qubits
+    # are usable (no holes, hanging strips, or disconnected qubits).
+    # Update this to use the shortest path length computed on the device's
+    # connectivity graph (ex using the output of Floyd-Warshall).
     return manhattan_dist(*gate_qubits) - manhattan_dist(*gate_after)
 
 
@@ -68,11 +69,6 @@ def __init__(self, initial_mapping: Dict[cirq.Qid, cirq.GridQubit] = {}):
         self.logical_to_physical = initial_mapping
         self.physical_to_logical = {v: k for k, v in initial_mapping.items()}
 
-    def insert(self, logical_q: cirq.Qid, physical_q: cirq.GridQubit) -> None:
-        """Inserts a mapping between a logical and physical qubit."""
-        self.logical_to_physical[logical_q] = physical_q
-        self.physical_to_logical[physical_q] = logical_q
-
     def swap_physical(self, q1: cirq.GridQubit, q2: cirq.GridQubit) -> None:
         """Updates the mapping by swapping two physical qubits."""
         logical_q1 = self.physical_to_logical.get(q1)
@@ -164,7 +160,10 @@ def _maximum_consecutive_positive_effect_impl(
         """
         total_cost = 0
         for gate in gates:
-            assert len(gate.qubits) <= 2
+            if len(gate.qubits) > 2:
+                raise ValueError(
+                    "Cannot compute maximum consecutive positive effect on gates with >2 qubits."
+                )
             if len(gate.qubits) != 2:
                 # Single-qubit gates would not be affected by the swap. We can
                 # treat the change in cost as 0 for those.
diff --git a/recirq/quantum_chess/mcpe_utils_test.py b/recirq/quantum_chess/mcpe_utils_test.py
index 6053e009..cd3a6e09 100644
--- a/recirq/quantum_chess/mcpe_utils_test.py
+++ b/recirq/quantum_chess/mcpe_utils_test.py
@@ -6,6 +6,43 @@
 import recirq.quantum_chess.mcpe_utils as mcpe
 
 
+def test_manhattan_distance():
+    assert mcpe.manhattan_dist(cirq.GridQubit(0, 0), cirq.GridQubit(0, 0)) == 0
+    assert mcpe.manhattan_dist(cirq.GridQubit(1, 2), cirq.GridQubit(1, 2)) == 0
+    assert mcpe.manhattan_dist(cirq.GridQubit(1, 2), cirq.GridQubit(3, 4)) == 4
+    assert mcpe.manhattan_dist(cirq.GridQubit(3, 4), cirq.GridQubit(1, 2)) == 4
+    assert mcpe.manhattan_dist(cirq.GridQubit(-1, 2), cirq.GridQubit(3,
+                                                                     -4)) == 10
+
+
+def test_swap_map_fn():
+    x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(3))
+    swap = mcpe.swap_map_fn(x, y)
+    assert swap(x) == y
+    assert swap(y) == x
+    assert swap(z) == z
+    assert cirq.Circuit(cirq.ISWAP(x, z), cirq.ISWAP(y, z), cirq.ISWAP(
+        x, y)).transform_qubits(swap) == cirq.Circuit(cirq.ISWAP(y, z),
+                                                      cirq.ISWAP(x, z),
+                                                      cirq.ISWAP(y, x))
+
+
+def test_effect_of_swap():
+    a1, a2, a3, b1, b2, b3 = cirq.GridQubit.rect(2, 3)
+    # If there's a gate operating on (a1, a3), then swapping a1 and a2 will
+    # bring the gate's qubits closer together by 1.
+    assert mcpe.effect_of_swap((a1, a2), (a1, a3)) == 1
+    # In reverse, a gate operating on (a2, a3) will get worse by 1 when swapping
+    # (a1, a2).
+    assert mcpe.effect_of_swap((a1, a2), (a2, a3)) == -1
+    # If the qubits to be swapped are completely independent of the gate's
+    # qubits, then there's no effect on the gate.
+    assert mcpe.effect_of_swap((a1, a2), (b1, b2)) == 0
+    # We can also measure the effect of swapping non-adjacent qubits (although
+    # we would never be able to do this with a real SWAP gate).
+    assert mcpe.effect_of_swap((a1, a3), (a1, b3)) == 2
+
+
 def test_peek():
     x, y, z = (cirq.NamedQubit(f'q{i}') for i in range(3))
     g = [cirq.ISWAP(x, y), cirq.ISWAP(x, z), cirq.ISWAP(y, z)]
diff --git a/recirq/quantum_chess/swap_updater.py b/recirq/quantum_chess/swap_updater.py
index 25bc50c5..9ba3669c 100644
--- a/recirq/quantum_chess/swap_updater.py
+++ b/recirq/quantum_chess/swap_updater.py
@@ -30,7 +30,9 @@ def _satisfies_adjacency(gate: cirq.Operation) -> bool:
 
     Requires the input to be either a binary or unary operation on GridQubits.
     """
-    assert len(gate.qubits) <= 2
+    if len(gate.qubits) > 2:
+        raise ValueError(
+            "Cannot determine physical adjacency for gates with > 2 qubits")
     if len(gate.qubits) < 2:
         return True
     q1, q2 = gate.qubits
@@ -40,10 +42,8 @@ def _satisfies_adjacency(gate: cirq.Operation) -> bool:
 def generate_decomposed_swap(
         q1: cirq.Qid, q2: cirq.Qid) -> Generator[cirq.Operation, None, None]:
     """Generates a SWAP operation using sqrt-iswap gates."""
-    yield cirq.ISWAP(q1, q2)**0.5
-    yield cirq.ISWAP(q1, q2)**0.5
-    yield cirq.ISWAP(q1, q2)**0.5
-    yield cirq.ISWAP(q1, q2)**0.5
+    yield from cirq.google.optimized_for_sycamore(
+        cirq.Circuit(cirq.SWAP(q1, q2))).all_operations()
 
 
 class SwapUpdater:
@@ -111,16 +111,19 @@ def update_iteration(self) -> Generator[cirq.Operation, None, None]:
                 # physical_gate needs to be fixed up with some swaps.
                 active_physical_gates.append(physical_gate)
 
-        # If all the active gates in this pass were added to the final circuit,
-        # then we have nothing left to do until we make another pass and get the
-        # newly active gates.
+        # If all the active gates in this pass were already optimal + added to
+        # the final circuit, then we have nothing left to do until we make
+        # another pass and get the newly active gates.
         if not active_physical_gates:
             return
 
         candidates = set(self.generate_candidate_swaps(active_physical_gates))
-        # This could happen on irregular topologies or when some device
-        # qubits are disallowed.
-        assert candidates, 'no swaps found that will improve the circuit'
+        if not candidates:
+            # This should never happen for reasonable initial mappings.
+            # For example, it can happen when the initial mapping placed a
+            # gate's qubits on disconnected components in the device
+            # connectivity graph.
+            raise ValueError("no swaps founds that will improve the circuit")
         chosen_swap = max(
             candidates,
             key=lambda swap: self.dlists.maximum_consecutive_positive_effect(
diff --git a/recirq/quantum_chess/swap_updater_test.py b/recirq/quantum_chess/swap_updater_test.py
index 8e0177e8..0d897d6c 100644
--- a/recirq/quantum_chess/swap_updater_test.py
+++ b/recirq/quantum_chess/swap_updater_test.py
@@ -94,11 +94,33 @@ def test_pentagonal_split_and_merge():
     updated_circuit = cirq.Circuit(
         SwapUpdater(circuit, grid_3x2, initial_mapping).add_swaps())
     for op in updated_circuit.all_operations():
-        assert len(op.qubits) == 2
-        q1, q2 = op.qubits
-        assert q1 in grid_3x2
-        assert q2 in grid_3x2
-        assert q1.is_adjacent(q2)
+        assert len(op.qubits) <= 2
+        assert all(q in grid_3x2 for q in op.qubits)
+        if len(op.qubits) == 2:
+            q1, q2 = op.qubits
+            assert q1.is_adjacent(q2)
+
+
+def test_already_optimal():
+    grid_2x3 = cirq.GridQubit.rect(2, 3)
+    logical_qubits = list(
+        cirq.NamedQubit(f'{x}{i}') for x in ('a', 'b') for i in range(3))
+    initial_mapping = dict(zip(logical_qubits, grid_2x3))
+    a1, a2, a3, b1, b2, b3 = logical_qubits
+    # Circuit has gates that already operate only on adjacent qubits.
+    circuit = cirq.Circuit(
+        cirq.ISWAP(a1, b1),
+        cirq.ISWAP(a2, a3),
+        cirq.ISWAP(b1, b2),
+        cirq.ISWAP(a3, b3),
+        cirq.ISWAP(b2, b3),
+    )
+    updated_circuit = cirq.Circuit(
+        SwapUpdater(circuit, grid_2x3, initial_mapping).add_swaps())
+    # The circuit was already optimal, so we don't add any extra operations,
+    # just map the logical qubits to physical qubits.
+    assert circuit.transform_qubits(
+        lambda q: initial_mapping.get(q)) == updated_circuit
 
 
 def test_decomposed_swaps():
@@ -109,9 +131,6 @@ def test_decomposed_swaps():
     # First iteration adds a swap between Q0 and Q1.
     # generate_decomposed_swap implements that swap operation as four
     # sqrt-iswaps.
-    assert list(updater.update_iteration()) == [
-        cirq.ISWAP(Q[0], Q[1])**0.5,
-        cirq.ISWAP(Q[0], Q[1])**0.5,
-        cirq.ISWAP(Q[0], Q[1])**0.5,
-        cirq.ISWAP(Q[0], Q[1])**0.5,
-    ]
+    assert cirq.Circuit(
+        updater.update_iteration()) == cirq.google.optimized_for_sycamore(
+            cirq.Circuit(cirq.SWAP(Q[0], Q[1])))

From dc50145bf0b6f65af1b5de1b7bcf9d75a6c1469c Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Fri, 19 Feb 2021 14:28:17 -0800
Subject: [PATCH 6/7] refactor swap updater tests

---
 recirq/quantum_chess/swap_updater_test.py | 52 +++++++++++++----------
 1 file changed, 30 insertions(+), 22 deletions(-)

diff --git a/recirq/quantum_chess/swap_updater_test.py b/recirq/quantum_chess/swap_updater_test.py
index 0d897d6c..0a4982b9 100644
--- a/recirq/quantum_chess/swap_updater_test.py
+++ b/recirq/quantum_chess/swap_updater_test.py
@@ -14,8 +14,7 @@
 #   Q0 = Q1 = Q2
 #   ||   ||   ||
 #   Q3 = Q4 = Q5
-FIGURE_9A_PHYSICAL_QUBITS = list(
-    cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+FIGURE_9A_PHYSICAL_QUBITS = cirq.GridQubit.rect(2, 3)
 # Circuit from figure 9a in
 # https://ieeexplore.ieee.org/abstract/document/8976109.
 FIGURE_9A_CIRCUIT = cirq.Circuit(cirq.CNOT(q[0], q[2]), cirq.CNOT(q[5], q[2]),
@@ -82,23 +81,25 @@ def test_example_9():
 
 
 def test_pentagonal_split_and_merge():
-    grid_3x2 = list(
-        cirq.GridQubit(row, col) for row in range(2) for col in range(3))
+    grid_2x3 = cirq.GridQubit.rect(2, 3, 4, 4)
     logical_qubits = list(
         cirq.NamedQubit(f'{x}{i}') for x in ('a', 'b') for i in range(3))
-    initial_mapping = dict(zip(logical_qubits, grid_3x2))
+    initial_mapping = dict(zip(logical_qubits, grid_2x3))
     a1, a2, a3, b1, b2, b3 = logical_qubits
-    circuit = cirq.Circuit(qm.normal_move(a1, b1), qm.normal_move(a2, a3),
-                           qm.merge_move(a3, b1, b3))
+    logical_circuit = cirq.Circuit(qm.normal_move(a1, b1),
+                                   qm.normal_move(a2, a3),
+                                   qm.merge_move(a3, b1, b3))
 
-    updated_circuit = cirq.Circuit(
-        SwapUpdater(circuit, grid_3x2, initial_mapping).add_swaps())
-    for op in updated_circuit.all_operations():
-        assert len(op.qubits) <= 2
-        assert all(q in grid_3x2 for q in op.qubits)
-        if len(op.qubits) == 2:
-            q1, q2 = op.qubits
-            assert q1.is_adjacent(q2)
+    updater = SwapUpdater(logical_circuit, grid_2x3, initial_mapping)
+    updated_circuit = cirq.Circuit(updater.add_swaps())
+
+    # Whereas the original circuit's initial mapping was not valid due to
+    # adjacency constraints, the updated circuit is valid.
+    device = cirq.google.Sycamore
+    with pytest.raises(ValueError):
+        device.validate_circuit(
+            logical_circuit.transform_qubits(lambda q: initial_mapping.get(q)))
+    device.validate_circuit(updated_circuit)
 
 
 def test_already_optimal():
@@ -117,8 +118,8 @@ def test_already_optimal():
     )
     updated_circuit = cirq.Circuit(
         SwapUpdater(circuit, grid_2x3, initial_mapping).add_swaps())
-    # The circuit was already optimal, so we don't add any extra operations,
-    # just map the logical qubits to physical qubits.
+    # The circuit was already optimal, so we didn't need to add any extra
+    # operations, just map the logical qubits to physical qubits.
     assert circuit.transform_qubits(
         lambda q: initial_mapping.get(q)) == updated_circuit
 
@@ -129,8 +130,15 @@ def test_decomposed_swaps():
     updater = SwapUpdater(FIGURE_9A_CIRCUIT, Q, initial_mapping,
                           generate_decomposed_swap)
     # First iteration adds a swap between Q0 and Q1.
-    # generate_decomposed_swap implements that swap operation as four
-    # sqrt-iswaps.
-    assert cirq.Circuit(
-        updater.update_iteration()) == cirq.google.optimized_for_sycamore(
-            cirq.Circuit(cirq.SWAP(Q[0], Q[1])))
+    # generate_decomposed_swap decomposes that into sqrt-iswap operations.
+    assert list(updater.update_iteration()) == list(
+        generate_decomposed_swap(Q[0], Q[1]))
+
+    # Whatever the decomposed operations are, they'd better be equivalent to a
+    # SWAP.
+    swap_unitary = cirq.unitary(cirq.Circuit(cirq.SWAP(Q[0], Q[1])))
+    generated_unitary = cirq.unitary(
+        cirq.Circuit(generate_decomposed_swap(Q[0], Q[1])))
+    cirq.testing.assert_allclose_up_to_global_phase(swap_unitary,
+                                                    generated_unitary,
+                                                    atol=1e-8)

From 770968d9bd4f5d6f9942e112c1c1661edb731e3a Mon Sep 17 00:00:00 2001
From: Jack Weinstein <jweinstein@google.com>
Date: Mon, 26 Apr 2021 11:53:00 -0700
Subject: [PATCH 7/7] update datadryad download URLs since they appear to have
 changed

---
 recirq/fermi_hubbard/publication.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/recirq/fermi_hubbard/publication.py b/recirq/fermi_hubbard/publication.py
index 89b50092..0b87ee39 100644
--- a/recirq/fermi_hubbard/publication.py
+++ b/recirq/fermi_hubbard/publication.py
@@ -241,10 +241,10 @@ def fetch_publication_data(
 
     base_url = "https://datadryad.org/stash/downloads/file_stream/"
     data = {
-        "gaussians_1u1d_nofloquet": "706210",
-        "gaussians_1u1d": "706211",
-        "trapping_2u2d": "706212",
-        "trapping_3u3d": "706213"
+        "gaussians_1u1d_nofloquet": "451326",
+        "gaussians_1u1d": "451327",
+        "trapping_2u2d": "451328",
+        "trapping_3u3d": "451329"
     }
     if exclude is not None:
         data = {path: key for path, key in data.items() if path not in exclude}