Add zip_longest support to the cirq_google sweep proto (#6815)

Author: BichengYing

cirq-google/cirq_google/api/v2/run_context.proto

@@ -85,6 +85,13 @@ message SweepFunction {
     // "a": 1.0, "b": 3.0
     // Note: if one sweep is shorter, the others will be truncated.
     ZIP = 2;
+
+    // A zip product of parameter sweeps with length as the longest one.
+    //
+    // Suppose we zip_longest([sweep.points(a, [1, 2]), sweep.points(b, [3])]),
+    // the iterator will produce: {a: 1, b: 3} and {a: 2, b: 3}.
+    // The shorter sweeps will be filled by repeating their last value.
+    ZIP_LONGEST = 3;
   }
 
   FunctionType function_type = 1;

cirq-google/cirq_google/api/v2/run_context_pb2.py

@@ -73,6 +73,10 @@ def sweep_to_proto(
         out.sweep_function.function_type = run_context_pb2.SweepFunction.PRODUCT
         for factor in sweep.factors:
             sweep_to_proto(factor, out=out.sweep_function.sweeps.add())
+    elif isinstance(sweep, cirq.ZipLongest):
+        out.sweep_function.function_type = run_context_pb2.SweepFunction.ZIP_LONGEST
+        for s in sweep.sweeps:
+            sweep_to_proto(s, out=out.sweep_function.sweeps.add())
     elif isinstance(sweep, cirq.Zip):
         out.sweep_function.function_type = run_context_pb2.SweepFunction.ZIP
         for s in sweep.sweeps:
@@ -129,6 +133,8 @@ def sweep_from_proto(msg: run_context_pb2.Sweep) -> cirq.Sweep:
             return cirq.Product(*factors)
         if func_type == run_context_pb2.SweepFunction.ZIP:
             return cirq.Zip(*factors)
+        if func_type == run_context_pb2.SweepFunction.ZIP_LONGEST:
+            return cirq.ZipLongest(*factors)
 
         raise ValueError(f'invalid sweep function type: {func_type}')
     if which == 'single_sweep':

cirq-google/cirq_google/api/v2/run_context_pb2.pyi

@@ -68,6 +68,7 @@ def _values(self) -> Iterator[float]:
                 + (cirq.Points('g', [1, 2]) * cirq.Points('h', [-1, 0, 1]))
             )
         ),
+        cirq.ZipLongest(cirq.Points('a', [1.0, 2.0, 3.0]), cirq.Points('b', [1.0])),
         # Sweep with constant. Type ignore is because cirq.Points type annotated with floats.
         cirq.Points('a', [None]),  # type: ignore[list-item]
         cirq.Points('a', [None]) * cirq.Points('b', [1, 2, 3]),  # type: ignore[list-item]