Add _decompose_with_context_ protocol to enable passing qubit manager within decompose

[tanujkhattar]

This PR fixes the first 3 tasks of #6040. Specifically, it does the following:

1. Adds a new abstraction for QubitManager and a trivial implementation called SimpleQubitManager which always allocates a new (internal qid type) CleanQubit / BorrowableQubit for qalloc / qfree requests.
2. Adds a new _decompose_with_context_ protocol which expects a new parameter context: cirq.DecompositionContext. cirq.DecompositionContext is a newly introduced dataclass that can be used to capture the arguments that can be passed around to gates as part of the decompose protocol; but in a type-safe way. It is a type-safe version of _decompose_with_options_(self, **kwargs) proposed in Quantum Memory Management for Cirq #6040 (comment) such that the same protocol can be used to support additional arguments, in a backwards compatible way, by simply extending the cirq.DecompositionContext dataclass.
3. Updates the cirq.decompose, cirq.decompose_once and cirq.decompose_once_with_qubits protocols to first try _decompose_with_context_ and then fallback on _decompose_ if former is not found / returns None or NotImplemented.
4. Updates existing composite gates; like ControlledGate, ControlledOperation, ClassicallyControlledGate, TaggedOperation etc. to implement the _decompose_with_context_ protocol and correctly pass around the context to decompose of sub-gates and operations.

Since cirq.decompose is a fairly complex component, I think it's important to see the different parts working together nicely in a single PR. Hence, I've chosen to open a single (slightly longer) PR instead of two. If there are strong opinions otherwise, I can also split up this PR into two smaller PRs - one that introduces QubitManager and other that adds the _decompose_with_context_ protocol.

cc @dstrain115 @daxfohl @NoureldinYosri

[NoureldinYosri]

looks good. adding a context to _try_decompose_into_operations_and_qubits is what is left

[NoureldinYosri]

_try_decompose_into_operations_and_qubits should also get a context. if it doesn't then cirq.unitary and cirq.apply_unitaries and other places that use it could fail when applied to operations that allocate ancillas. for now they work because the decompositions we have use global context which will be gone soon.

once this PR is merged we will need to update those as well.

[tanujkhattar]

Yeah, for _try_decompose_into_operations_and_qubits; I plan to always use SimpleQubitManager as the default context instead of expecting a context in the input. This is because the internal method is mostly used by protocols which should correctly work irrespective of the qubit allocation strategy ?

I do, however, plan to modify the return type of _try_decompose_into_operations_and_qubits in a follow-up PR and update all the call-sites so that the logic to identify work qubits and ancilla qubits happens within _try_decompose_into_operations_and_qubits.

[NoureldinYosri]

they protocols won't care about the allocation strategy. but the qubit names will. if one of those protocols has to decompose twice then it could end up creating a gate acting on the same qubit twice. e.g.

class G1:
   def _decompose_with_context_(self, qubits, context):
       ancilla = context.qm.qalloc(1)
       yield CNOT(qubits[0], ancillas[0])
       ...

class G2:
   def _decompose_with_context_(self, qubits, context):
       ancilla = context.qm.qalloc(1)
       yield G1(*(ancillas + qubits))

if we always create a new qubit manager then when decomposing G2() the CNOT will endup having its two inputs being the same qubit => raising an exception

the context here won't be about the allocation strategy but about book keeping of ancilla names

---

An alternative would be to rename the input qubits but that would be hacky and error prone

[tanujkhattar]

Hmm, maybe we should change the implementation of SimpleQubitManager to add a prefix by default which is unique for each instantiation of the SimpleQubitManager (eg: using the current epoch time) ?

The reason I am hesitant to add a context argument to _try_decompose_into_operations_and_qubits is because multiple calls to decompose_once() is not the same as a single call cirq.decompose(op, keep=predicate) when we have a qubit manager around (this goes back to the dfs vs bfs traversal; cirq.decompose() is dfs but multiple calls to cirq.decompose_once() would be bfs).

So, we can either enforce the use of a custom qubit manager that is guaranteed to always allocate new qubits (eg: a modified version of SimpleQubitManager) or the call sites become more complicated and need to have this check always.

I say we punt the issue of updating _try_decompose_into_operations_and_qubits and corresponding call sites within other protocols to a follow-up PR and focus on the rest of the changes in this PR for now. I'll send a follow-up soon.

[NoureldinYosri]

sg. the prefix shouldn't be added by default though, when a user uses the SimpleQubitManager they won't want to see the extra information which will be confusing.

Anyway, yes lets leave it to a follow up PR

[NoureldinYosri]

LGTM