str/repr blows up for large disjoint states

[daxfohl]

Description of the issue

The repr and str functions of TrialResult and subclasses all compute the full state of the system as a kron of all disjoint states before displaying. In large disjoint states, this causes the system to hang. This can crop up in unusual places like assertions and stack traces. See #4360 (comment).

How to reproduce the issue

In the density_matrix_simulator_test::test_large_untangled_okay unit test, change the number of qubits to mismatch. The simulation works fine, but the test then hangs at assert set(result._final_step_result.qubits) == set(cirq.LineQubit.range(59)), because it tries to calculate the repr of result to display. This led to users thinking that the simulation itself was hanging, which was not the case.

I propose we update these to display the individual state spaces and associated qubits, rather than joining into a single space prior to output.

[dstrain115]

Discussion from cirq cync: This seems bad. We should fix this fairly urgently, but we don't have a solution yet or someone assigned.

[daxfohl]

I can take this. It's related to the state separation functionality introduced in #4100. Updating the code to output these as a list of independent state vectors or density matrices should be straightforward.

I think as part of this we should also add a public function to query for the independent state space that contains a qubit too. (e.g. trial_result.get_state(q1).qubits and trial_result.get_state(q1).density_matrix would return the qubits and density matrix for whatever state space q1 is in). Currently that's all private as an implementation detail, but if we're updating the str/repr to display this information, then it seems logical that we should have an easy way for users to get that raw data too. @95-martin-orion @MichaelBroughton thoughts?

[95-martin-orion]

I think as part of this we should also add a public function to query for the independent state space that contains a qubit too.

This sounds like a reasonable, low-cost addition to me - though I would prefer get_state_containing_qubit to be clear about what the method is doing. (e.g. to avoid confusion that a state is only associated with the provided qubits)

[daxfohl]

Updated linked PR with this function.