def rewrite_program(raw_prog: Program, qecc: QECC) -> Program:
if qecc.k != 1:
raise UnsupportedQECCError("code must have k = 1")
if raw_prog.defined_gates:
raise UnsupportedProgramError("does not support DEFGATE")
# Assign indices to qubit placeholders in the raw program.
raw_prog = quil.address_qubits(raw_prog)
new_prog = Program()
logical_qubits = {
index: new_logical_qubit(new_prog, qecc,
"logical_qubit_{}".format(index))
for index in raw_prog.get_qubits(indices=True)
}
# Construct ancilla code blocks.
ancilla_1 = new_logical_qubit(new_prog, qecc, "ancilla_1")
ancilla_2 = new_logical_qubit(new_prog, qecc, "ancilla_2")
# Classical scratch BIT registers for gates/measurements.
scratch_size = max(qecc.n, qecc.measure_scratch_size)
raw_scratch = new_prog.declare('scratch', 'BIT', scratch_size)
scratch = MemoryChunk(raw_scratch, 0, raw_scratch.declared_size)
_initialize_memory(new_prog, raw_scratch,
ancilla_1.qubits + ancilla_2.qubits)
# Classical scratch INTEGER registers.
raw_scratch_int = new_prog.declare('scratch_int', 'INTEGER', 2)
scratch_int = MemoryChunk(raw_scratch_int, 0,
raw_scratch_int.declared_size)
_initialize_memory(new_prog, raw_scratch_int,
ancilla_1.qubits + ancilla_2.qubits)
perform_error_correction = _make_error_corrector(new_prog, qecc, ancilla_1,
ancilla_2)
# Reset all logical qubits.
for block in logical_qubits.values():
qecc.encode_zero(new_prog, block, ancilla_1, scratch)
for inst in raw_prog.instructions:
if isinstance(inst, Gate):
gate_qubits = [
logical_qubits[index] for index in _gate_qubits(inst)
]
qecc.apply_gate(new_prog, inst.name, *gate_qubits)
# Perform error correction after every logical gate.
perform_error_correction(logical_qubits.values())
elif isinstance(inst, Measurement):
qubit = logical_qubits[_extract_qubit_index(inst.qubit)]
# This should really use its own ancilla instead of sharing with the error correction,
# but we need be extremely conservative with the number of qubits.
for _ in qecc.measure(new_prog, qubit, 0, inst.classical_reg,
ancilla_1, ancilla_2, scratch, scratch_int):
# Since measurements are taken multiple times for redundancy, we need to perform
# rounds of error correction during the measurement routine.
perform_error_correction(logical_qubits.values())
elif isinstance(inst, ResetQubit):
raise NotImplementedError(
"this instruction is not in the Quil spec")
elif isinstance(inst, JumpTarget):
new_prog.inst(JumpTarget(_mangle_label(inst.label)))
elif isinstance(inst, JumpConditional):
new_prog.inst(
type(inst)(_mangle_label(inst.target), inst.condition))
elif isinstance(inst, Jump):
new_prog.inst(Jump(_mangle_label(inst.target)))
elif isinstance(inst, Halt):
new_prog.append(inst)
elif isinstance(inst, Wait):
raise NotImplementedError()
elif isinstance(inst, Reset):
for block in logical_qubits.values():
qecc.encode_zero(new_prog, block.qubits, ancilla_1, scratch)
elif isinstance(inst, Declare):
new_prog.inst(inst)
elif isinstance(inst, Pragma):
new_prog.inst(inst)
elif any(
isinstance(inst, ClassicalInst)
for ClassicalInst in CLASSICAL_INSTRUCTIONS):
new_prog.inst(inst)
else:
raise UnsupportedProgramError("unsupported instruction: {}", inst)
return quil.address_qubits(new_prog)
