pass on address analysis interpreter info instead of sifting through lattice elements

Author: johnzl-777

src/bloqade/gemini/analysis/measurement_validation/analysis.py

@@ -13,8 +13,8 @@
 class _GeminiTerminalMeasurementValidationAnalysis(Forward[EmptyLattice]):
     keys = ("gemini.validate.terminal_measurement",)
 
-    address_analysis_results: ForwardFrame
     measurement_analysis_results: ForwardFrame
+    unique_qubits_allocated: int = 0
     terminal_measurement_encountered: bool = False
     lattice = EmptyLattice
 
@@ -35,28 +35,30 @@ def name(self) -> str:
         return "Gemini Terminal Measurement Validation"
 
     def get_required_analyses(self) -> list[type]:
-        return [measure_id.MeasurementIDAnalysis, address.AddressAnalysis]
+        return [measure_id.MeasurementIDAnalysis]
 
     def set_analysis_cache(self, cache: dict[type, Any]) -> None:
         self.analysis_cache.update(cache)
 
     def run(self, method: ir.Method) -> tuple[Any, list[ir.ValidationError]]:
 
         # get the data out of the cache and forward it to the underlying analysis
-        address_analysis_results = self.analysis_cache.get(address.AddressAnalysis)
         measurement_analysis_results = self.analysis_cache.get(
             measure_id.MeasurementIDAnalysis
         )
 
-        assert (
-            address_analysis_results is not None
-        ), "Address analysis results not found in cache"
+        address_analysis = address.AddressAnalysis(dialects=method.dialects)
+        address_analysis.run(method)
+        unique_qubits_allocated = address_analysis.qubit_count
+
         assert (
             measurement_analysis_results is not None
         ), "Measurement ID analysis results not found in cache"
 
         analysis = _GeminiTerminalMeasurementValidationAnalysis(
-            method.dialects, address_analysis_results, measurement_analysis_results
+            method.dialects,
+            measurement_analysis_results,
+            unique_qubits_allocated=unique_qubits_allocated,
         )
 
         frame, _ = analysis.run(method)

src/bloqade/gemini/analysis/measurement_validation/impls.py

@@ -4,7 +4,6 @@
 
 from bloqade import qubit, gemini
 from bloqade.analysis.address.impls import Func as AddressFuncMethodTable
-from bloqade.analysis.address.lattice import AddressReg, AddressQubit
 from bloqade.analysis.measure_id.lattice import MeasureIdTuple
 
 from .analysis import _GeminiTerminalMeasurementValidationAnalysis
@@ -60,30 +59,8 @@ def terminal_measure(
             )
             return (interp.lattice.bottom(),)
 
-        # If we confirm there isn't a duplicate terminal measurement,
-        # now we need to check that for all the qubits that were spawned,
-        # they are all consumed by this measurement
-        address_analysis_results = interp.address_analysis_results
         measurement_analysis_results = interp.measurement_analysis_results
-
-        # Figure out the total number of qubits spawned, keeping in mind that if a user
-        # "shuffles" the qubits (puts them in a new container, splits one off from a container type, etc.)
-        # it should be accounted for. This would be much cleaner if there was a way to propagate the
-        # final qubit count saved in the actual interpreter for address analysis...
-        witnessed_qubits = set()
-        total_qubits_allocated = 0
-        for address_lattice_elem in address_analysis_results.entries.values():
-
-            match address_lattice_elem:
-                case AddressReg(data=data):
-                    for data_elem in data:
-                        if data_elem not in witnessed_qubits:
-                            witnessed_qubits.add(data_elem)
-                            total_qubits_allocated += 1
-                case AddressQubit(data=data):
-                    if data not in witnessed_qubits:
-                        witnessed_qubits.add(data)
-                        total_qubits_allocated += 1
+        total_qubits_allocated = interp.unique_qubits_allocated
 
         # could make these proper exceptions but would be tricky to communicate to user
         # without revealing under-the-hood details