1# Copyright 2018 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from typing import cast, Sequence, TYPE_CHECKING
16
17from cirq import devices, ops, protocols
18from cirq.contrib.acquaintance.permutation import PermutationGate, update_mapping
19
20if TYPE_CHECKING:
21    import cirq
22
23
24def assert_permutation_decomposition_equivalence(gate: PermutationGate, n_qubits: int) -> None:
25    qubits = devices.LineQubit.range(n_qubits)
26    operations = protocols.decompose_once_with_qubits(gate, qubits)
27    operations = list(cast(Sequence['cirq.Operation'], ops.flatten_op_tree(operations)))
28    mapping = {cast(ops.Qid, q): i for i, q in enumerate(qubits)}
29    update_mapping(mapping, operations)
30    expected_mapping = {qubits[j]: i for i, j in gate.permutation().items()}
31    assert mapping == expected_mapping, (
32        "{!r}.permutation({}) doesn't match decomposition.\n"
33        '\n'
34        'Actual mapping:\n'
35        '{}\n'
36        '\n'
37        'Expected mapping:\n'
38        '{}\n'.format(
39            gate, n_qubits, [mapping[q] for q in qubits], [expected_mapping[q] for q in qubits]
40        )
41    )
42