quil_rs/program/
defgate_sequence_expansion.rs

1/// This module implements the expansion of sequence gate definitions in a Quil program,
2/// as well as the associated source map that tracks how the original
3/// instructions map to the expanded instructions.
4use std::{collections::HashMap, ops::Range, vec};
5
6use indexmap::{IndexMap, IndexSet};
7
8use crate::{
9    expression::Expression,
10    instruction::{
11        DefGateSequenceExpansionError, GateDefinition, GateSignature, GateSpecification,
12        Instruction,
13    },
14    program::{InstructionIndex, SourceMap, SourceMapEntry},
15};
16
17use super::source_map::{ExpansionResult, SourceMapIndexable};
18
19/// Details about the expansion of a sequence gate definition
20#[derive(Clone, Debug, PartialEq)]
21pub struct DefGateSequenceExpansion<'a> {
22    /// The signature of the sequence gate definition that was
23    /// used to expand the instruction.
24    ///
25    /// Note, technically, the gate name itself is sufficient to identify
26    /// the sequence gate definition, since gate names are unique within
27    /// a program. Nevertheless, we include the full signature here
28    /// for later reference within error messages.
29    source_signature: crate::instruction::GateSignature<'a>,
30
31    /// The target instruction indices produced by the expansion
32    range: Range<InstructionIndex>,
33
34    /// Sequence gate definitions may refer to other sequence gate definitions
35    /// per the Quil specification. As such, we need to track how the first-level
36    /// sequence instructions map to nested sequence gate definition expansion.
37    nested_expansions: SourceMap<InstructionIndex, ExpansionResult<DefGateSequenceExpansion<'a>>>,
38}
39
40impl<'a> DefGateSequenceExpansion<'a> {
41    /// Borrow the source gate signature of the sequence gate definition
42    pub(crate) fn source_signature(&self) -> &GateSignature<'a> {
43        &self.source_signature
44    }
45
46    /// Returns the range of target instruction indices produced by the expansion
47    pub fn range(&self) -> &Range<InstructionIndex> {
48        &self.range
49    }
50
51    /// Returns the nested expansions of this sequence gate definition
52    pub fn nested_expansions(
53        &self,
54    ) -> &SourceMap<InstructionIndex, ExpansionResult<DefGateSequenceExpansion<'a>>> {
55        &self.nested_expansions
56    }
57}
58
59impl SourceMapIndexable<InstructionIndex> for DefGateSequenceExpansion<'_> {
60    fn contains(&self, other: &InstructionIndex) -> bool {
61        self.range.contains(other)
62    }
63}
64
65impl<'a> SourceMapIndexable<GateSignature<'a>> for DefGateSequenceExpansion<'a> {
66    fn contains(&self, other: &GateSignature) -> bool {
67        &self.source_signature == other
68    }
69}
70
71type SequenceGateDefinitionSourceMap<'a> =
72    SourceMap<InstructionIndex, ExpansionResult<DefGateSequenceExpansion<'a>>>;
73
74/// A utility to expand sequence gate definitions in a Quil program.
75#[derive(Clone, Debug, PartialEq)]
76pub(crate) struct ProgramDefGateSequenceExpander<'a, F> {
77    gate_definitions: &'a IndexMap<String, GateDefinition>,
78    filter: F,
79}
80
81#[derive(Clone, Debug, PartialEq)]
82pub(crate) struct ExpandedInstructionsWithSourceMap<'a> {
83    pub(crate) instructions: Vec<Instruction>,
84    pub(crate) source_map: SequenceGateDefinitionSourceMap<'a>,
85}
86
87struct ExpansionStack(IndexSet<String>);
88
89impl ExpansionStack {
90    fn new() -> Self {
91        Self(IndexSet::new())
92    }
93
94    /// Check if the name is in the stack and, if so, return an error.
95    fn check(&self, name: impl AsRef<str>) -> Result<(), DefGateSequenceExpansionError> {
96        if self.0.contains(name.as_ref()) {
97            let cycle = self.0.iter().cloned().collect();
98            Err(DefGateSequenceExpansionError::CyclicSequenceGateDefinition(
99                cycle,
100            ))
101        } else {
102            Ok(())
103        }
104    }
105
106    /// Execute a closure with an gate added to the stack.
107    fn with_gate_sequence<F, R>(&mut self, name: String, f: F) -> R
108    where
109        F: FnOnce(&mut Self) -> R,
110    {
111        let must_pop = self.0.insert(name);
112        let result = f(self);
113        if must_pop {
114            self.0.pop();
115        }
116        result
117    }
118}
119
120impl<'a, F> ProgramDefGateSequenceExpander<'a, F>
121where
122    F: Fn(&str) -> bool,
123{
124    /// Creates a new `ProgramDefGateSequenceExpander`.
125    ///
126    /// # Arguments
127    ///
128    /// * `gate_definitions` - A reference to the gate definitions of the program.
129    /// * `filter` - A filter to apply to the gate definitions, allowing for selective
130    ///   expansion.
131    pub(crate) fn new(gate_definitions: &'a IndexMap<String, GateDefinition>, filter: F) -> Self {
132        Self {
133            gate_definitions,
134            filter,
135        }
136    }
137
138    /// Expands sequence gate definitions in the provided instructions.
139    pub(crate) fn expand(
140        &self,
141        source_instructions: &[Instruction],
142    ) -> Result<Vec<Instruction>, DefGateSequenceExpansionError> {
143        self.expand_without_source_map_impl(source_instructions, &mut ExpansionStack::new())
144    }
145
146    /// Expands sequence gate definitions in the provided instructions and returns a source map
147    /// detailing the expansion.
148    pub(crate) fn expand_with_source_map(
149        &self,
150        source_instructions: &'a [Instruction],
151    ) -> Result<ExpandedInstructionsWithSourceMap<'a>, DefGateSequenceExpansionError> {
152        let mut source_map = SourceMap::default();
153        self.expand_with_source_map_impl(
154            source_instructions,
155            &mut source_map,
156            &mut ExpansionStack::new(),
157        )
158        .map(|instructions| ExpandedInstructionsWithSourceMap {
159            instructions,
160            source_map,
161        })
162    }
163
164    fn expand_with_source_map_impl(
165        &self,
166        source_instructions: &[Instruction],
167        source_map: &mut SequenceGateDefinitionSourceMap<'a>,
168        stack: &mut ExpansionStack,
169    ) -> Result<Vec<Instruction>, DefGateSequenceExpansionError> {
170        let mut target_instructions = vec![];
171        for (source_instruction_index, source_instruction) in source_instructions.iter().enumerate()
172        {
173            if let Some((target_gate_instructions, gate_sequence_signature)) =
174                self.gate_sequence_from_instruction(source_instruction, stack)?
175            {
176                // If this instruction is a sequence gate definition, we need to expand it. Before
177                // doing so, we add the gate sequence signature to the `gate_expansion_stack`,
178                // so all nested expansions within this sequence have access to the stack of
179                // already expanded gate definitions.
180                let mut nested_expansions = SourceMap::default();
181                let recursive_target_gate_instructions = stack.with_gate_sequence(
182                    gate_sequence_signature.name().to_string(),
183                    |stack| {
184                        self.expand_with_source_map_impl(
185                            &target_gate_instructions,
186                            &mut nested_expansions,
187                            stack,
188                        )
189                    },
190                )?;
191
192                let target_instruction_start_index = InstructionIndex(target_instructions.len());
193                let target_instruction_end_index = InstructionIndex(
194                    target_instruction_start_index.0 + recursive_target_gate_instructions.len(),
195                );
196                source_map.entries.push(SourceMapEntry {
197                    source_location: InstructionIndex(source_instruction_index),
198                    target_location: ExpansionResult::Rewritten(DefGateSequenceExpansion {
199                        source_signature: gate_sequence_signature,
200                        range: target_instruction_start_index..target_instruction_end_index,
201                        nested_expansions,
202                    }),
203                });
204                target_instructions.extend(recursive_target_gate_instructions);
205            } else {
206                target_instructions.push(source_instruction.clone());
207                source_map.entries.push(SourceMapEntry {
208                    source_location: InstructionIndex(source_instruction_index),
209                    target_location: ExpansionResult::Unmodified(InstructionIndex(
210                        target_instructions.len() - 1,
211                    )),
212                });
213            }
214        }
215        Ok(target_instructions)
216    }
217
218    fn expand_without_source_map_impl(
219        &self,
220        source_instructions: &[Instruction],
221        stack: &mut ExpansionStack,
222    ) -> Result<Vec<Instruction>, DefGateSequenceExpansionError> {
223        let mut target_instructions = vec![];
224        for source_instruction in source_instructions {
225            if let Some((target_gate_instructions, source)) =
226                self.gate_sequence_from_instruction(source_instruction, stack)?
227            {
228                // If this instruction is a sequence gate definition, we need to expand it. Before
229                // doing so, we add the gate sequence signature to the the `gate_expansion_stack`,
230                // so all nested expansions within this sequence have access to the stack of
231                // already expanded gate definitions.
232                let recursive_target_gate_instructions = stack
233                    .with_gate_sequence(source.name().to_string(), |stack| {
234                        self.expand_without_source_map_impl(&target_gate_instructions, stack)
235                    })?;
236                target_instructions.extend(recursive_target_gate_instructions);
237            } else {
238                target_instructions.push(source_instruction.clone());
239            }
240        }
241        Ok(target_instructions)
242    }
243
244    /// Given an instruction, this function checks if it is a gate instruction that
245    /// matches a sequence gate definition. If it does and the gate name is included
246    /// by the [`ProgramDefGateSequenceExpander::filter`], it expands the gate
247    /// definition into a sequence of instructions and returns them along with the
248    /// signature of the gate definition.
249    ///
250    /// This also checks the `gate_expansion_stack` set to prevent cyclic expansions.
251    fn gate_sequence_from_instruction(
252        &self,
253        instruction: &Instruction,
254        stack: &ExpansionStack,
255    ) -> Result<Option<(Vec<Instruction>, GateSignature<'a>)>, DefGateSequenceExpansionError> {
256        if let Instruction::Gate(gate) = instruction {
257            if let Some(gate_definition) = self.gate_definitions.get(&gate.name) {
258                if let GateSpecification::Sequence(gate_sequence) = &gate_definition.specification {
259                    if (self.filter)(&gate.name) {
260                        if gate_definition.parameters.len() != gate.parameters.len() {
261                            return Err(DefGateSequenceExpansionError::ParameterCount {
262                                expected: gate_definition.parameters.len(),
263                                found: gate.parameters.len(),
264                            });
265                        }
266                        let gate_parameter_arguments = gate_definition
267                            .parameters
268                            .iter()
269                            .cloned()
270                            .zip(gate.parameters.iter().cloned())
271                            .collect::<HashMap<String, Expression>>();
272
273                        if !gate.modifiers.is_empty() {
274                            return Err(DefGateSequenceExpansionError::GateModifiersUnsupported(
275                                gate.modifiers.clone(),
276                            ));
277                        }
278                        let source = gate_definition.signature();
279                        stack.check(source.name())?;
280
281                        let target_gate_instructions = gate_sequence
282                            .expand(gate_parameter_arguments, gate.qubits.clone())?
283                            .into_iter()
284                            .map(Instruction::Gate)
285                            .collect::<Vec<_>>();
286
287                        return Ok(Some((target_gate_instructions, source)));
288                    }
289                }
290            }
291        }
292        Ok(None)
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use std::str::FromStr;
299
300    use crate::{instruction::GateSignature, Program};
301
302    use super::*;
303    use rstest::*;
304
305    /// A test case for the [`ProgramDefGateSequenceExpander`] functionality.
306    struct DefGateSequenceExpansionTestCase {
307        program: &'static str,
308        filter: Box<dyn Fn(&str) -> bool>,
309        expected: Result<&'static str, DefGateSequenceExpansionError>,
310        source_map_entry_builders:
311            Vec<SourceMapEntry<InstructionIndex, ExpansionResult<DefGateSequenceExpansionBuilder>>>,
312    }
313
314    /// Below we define a set of test cases for the `DefGateSequenceExpansion` functionality. We
315    /// cover valid and invalid expansions.
316    ///
317    /// Note, the error coverage is comprehensive with the exception of
318    /// [`DefGateSequenceExpansionError::UndefinedGateSequenceElementQubit`] and
319    /// [`DefGateSequenceExpansionError::InvalidGateSequenceElementQubit`], which
320    /// cover Quil errors that are impossible to parse or construct.
321    impl DefGateSequenceExpansionTestCase {
322        fn to_source_map(
323            &self,
324        ) -> SourceMap<InstructionIndex, ExpansionResult<DefGateSequenceExpansion<'_>>> {
325            SourceMap {
326                entries: self
327                    .source_map_entry_builders
328                    .iter()
329                    .map(|builder| SourceMapEntry {
330                        source_location: builder.source_location,
331                        target_location: match &builder.target_location {
332                            ExpansionResult::Rewritten(expansion) => expansion.build(),
333                            ExpansionResult::Unmodified(index) => {
334                                ExpansionResult::Unmodified(*index)
335                            }
336                        },
337                    })
338                    .collect(),
339            }
340        }
341
342        fn simple_1q_expansions() -> Self {
343            const QUIL: &str = r"
344DEFGATE seq2(%param1, %param2) a b AS SEQUENCE:
345    seq1(%param1) a
346    seq1(%param2) b
347
348DEFGATE seq1(%param1) a AS SEQUENCE:
349    RZ(%param1) a
350    RX(pi/2) a
351    RZ(%param1) a
352
353seq2(pi, pi/2) 0 1
354";
355            const EXPECTED_QUIL: &str = r"
356RZ(pi) 0
357RX(pi/2) 0
358RZ(pi) 0
359RZ(pi/2) 1
360RX(pi/2) 1
361RZ(pi/2) 1
362";
363            let source_map_entry_builders = vec![build_source_map_entry(
364                0,
365                DefGateSequenceExpansionBuilder::new(
366                    "seq2",
367                    &["param1", "param2"],
368                    &["a", "b"],
369                    0..6,
370                    vec![
371                        build_source_map_entry(
372                            0,
373                            DefGateSequenceExpansionBuilder::new(
374                                "seq1",
375                                &["param1"],
376                                &["a"],
377                                0..3,
378                                vec![
379                                    build_source_map_entry_copy(0, 0),
380                                    build_source_map_entry_copy(1, 1),
381                                    build_source_map_entry_copy(2, 2),
382                                ],
383                            ),
384                        ),
385                        build_source_map_entry(
386                            1,
387                            DefGateSequenceExpansionBuilder::new(
388                                "seq1",
389                                &["param1"],
390                                &["a"],
391                                3..6,
392                                vec![
393                                    build_source_map_entry_copy(0, 0),
394                                    build_source_map_entry_copy(1, 1),
395                                    build_source_map_entry_copy(2, 2),
396                                ],
397                            ),
398                        ),
399                    ],
400                ),
401            )];
402            Self {
403                program: QUIL,
404                filter: Box::new(|_| true),
405                expected: Ok(EXPECTED_QUIL),
406                source_map_entry_builders,
407            }
408        }
409
410        #[expect(clippy::too_many_lines)]
411        fn triple_recursize() -> Self {
412            const QUIL: &str = r"
413DEFGATE some_u2_cycle(%param1, %param2, %param3, %param4, %param5, %param6) a b AS SEQUENCE:
414    pmw3(%param1, %param2, %param3) a
415    pmw3(%param4, %param5, %param6) b
416
417DEFGATE pmw3(%param1, %param2, %param3) a AS SEQUENCE:
418    pmw(%param1) a
419    pmw(%param2) a
420    pmw(%param3) a
421
422DEFGATE pmw(%param1) a AS SEQUENCE:
423    RZ(%param1) a
424    RX(pi/2) a
425    RZ(-%param1) a
426
427some_u2_cycle(-pi, -pi/2, -pi/4, pi/4, pi/2, pi) 0 1
428";
429            const EXPECTED_QUIL: &str = r"
430RZ(-pi) 0
431RX(pi/2) 0
432RZ(-(-pi)) 0
433RZ(-pi/2) 0
434RX(pi/2) 0
435RZ(-(-pi/2)) 0
436RZ(-pi/4) 0
437RX(pi/2) 0
438RZ(-(-pi/4)) 0
439
440RZ(pi/4) 1
441RX(pi/2) 1
442RZ(-(pi/4)) 1
443RZ(pi/2) 1
444RX(pi/2) 1
445RZ(-(pi/2)) 1
446RZ(pi) 1
447RX(pi/2) 1
448RZ(-(pi)) 1
449";
450            let source_map_entry_builders = vec![build_source_map_entry(
451                0,
452                DefGateSequenceExpansionBuilder::new(
453                    "some_u2_cycle",
454                    &["param1", "param2", "param3", "param4", "param5", "param6"],
455                    &["a", "b"],
456                    0..18,
457                    vec![
458                        build_source_map_entry(
459                            0,
460                            DefGateSequenceExpansionBuilder::new(
461                                "pmw3",
462                                &["param1", "param2", "param3"],
463                                &["a"],
464                                0..9,
465                                vec![
466                                    build_source_map_entry(
467                                        0,
468                                        DefGateSequenceExpansionBuilder::new(
469                                            "pmw",
470                                            &["param1"],
471                                            &["a"],
472                                            0..3,
473                                            vec![
474                                                build_source_map_entry_copy(0, 0),
475                                                build_source_map_entry_copy(1, 1),
476                                                build_source_map_entry_copy(2, 2),
477                                            ],
478                                        ),
479                                    ),
480                                    build_source_map_entry(
481                                        1,
482                                        DefGateSequenceExpansionBuilder::new(
483                                            "pmw",
484                                            &["param1"],
485                                            &["a"],
486                                            3..6,
487                                            vec![
488                                                build_source_map_entry_copy(0, 0),
489                                                build_source_map_entry_copy(1, 1),
490                                                build_source_map_entry_copy(2, 2),
491                                            ],
492                                        ),
493                                    ),
494                                    build_source_map_entry(
495                                        2,
496                                        DefGateSequenceExpansionBuilder::new(
497                                            "pmw",
498                                            &["param1"],
499                                            &["a"],
500                                            6..9,
501                                            vec![
502                                                build_source_map_entry_copy(0, 0),
503                                                build_source_map_entry_copy(1, 1),
504                                                build_source_map_entry_copy(2, 2),
505                                            ],
506                                        ),
507                                    ),
508                                ],
509                            ),
510                        ),
511                        build_source_map_entry(
512                            1,
513                            DefGateSequenceExpansionBuilder::new(
514                                "pmw3",
515                                &["param1", "param2", "param3"],
516                                &["a"],
517                                9..18,
518                                vec![
519                                    build_source_map_entry(
520                                        0,
521                                        DefGateSequenceExpansionBuilder::new(
522                                            "pmw",
523                                            &["param1"],
524                                            &["a"],
525                                            0..3,
526                                            vec![
527                                                build_source_map_entry_copy(0, 0),
528                                                build_source_map_entry_copy(1, 1),
529                                                build_source_map_entry_copy(2, 2),
530                                            ],
531                                        ),
532                                    ),
533                                    build_source_map_entry(
534                                        1,
535                                        DefGateSequenceExpansionBuilder::new(
536                                            "pmw",
537                                            &["param1"],
538                                            &["a"],
539                                            3..6,
540                                            vec![
541                                                build_source_map_entry_copy(0, 0),
542                                                build_source_map_entry_copy(1, 1),
543                                                build_source_map_entry_copy(2, 2),
544                                            ],
545                                        ),
546                                    ),
547                                    build_source_map_entry(
548                                        2,
549                                        DefGateSequenceExpansionBuilder::new(
550                                            "pmw",
551                                            &["param1"],
552                                            &["a"],
553                                            6..9,
554                                            vec![
555                                                build_source_map_entry_copy(0, 0),
556                                                build_source_map_entry_copy(1, 1),
557                                                build_source_map_entry_copy(2, 2),
558                                            ],
559                                        ),
560                                    ),
561                                ],
562                            ),
563                        ),
564                    ],
565                ),
566            )];
567            Self {
568                program: QUIL,
569                filter: Box::new(|_| true),
570                expected: Ok(EXPECTED_QUIL),
571                source_map_entry_builders,
572            }
573        }
574
575        /// Test a program expansion where some instructions are not expanded
576        fn unexpanded_instructions() -> Self {
577            const QUIL: &str = r"
578DEFGATE seq2(%param1, %param2) a b AS SEQUENCE:
579    X a
580    seq1(%param2) b
581    H b
582    ISWAP a b
583
584DEFGATE seq1(%param1) a AS SEQUENCE:
585    RZ(%param1) a
586    RX(pi/2) a
587    RZ(%param1) a
588
589ISWAP 0 1
590seq2(pi, pi/2) 0 1
591MEASURE 0 ro[0]
592MEASURE 1 ro[1]
593";
594            const EXPECTED_QUIL: &str = r"
595ISWAP 0 1
596X 0
597RZ(pi/2) 1
598RX(pi/2) 1
599RZ(pi/2) 1
600H 1
601ISWAP 0 1
602MEASURE 0 ro[0]
603MEASURE 1 ro[1]
604";
605            let source_map_entry_builders = vec![
606                build_source_map_entry(0, ExpansionResult::Unmodified(InstructionIndex(0))),
607                build_source_map_entry(
608                    1,
609                    DefGateSequenceExpansionBuilder::new(
610                        "seq2",
611                        &["param1", "param2"],
612                        &["a", "b"],
613                        1..7,
614                        vec![
615                            build_source_map_entry_copy(0, 0),
616                            build_source_map_entry(
617                                1,
618                                DefGateSequenceExpansionBuilder::new(
619                                    "seq1",
620                                    &["param1"],
621                                    &["a"],
622                                    1..4,
623                                    vec![
624                                        build_source_map_entry_copy(0, 0),
625                                        build_source_map_entry_copy(1, 1),
626                                        build_source_map_entry_copy(2, 2),
627                                    ],
628                                ),
629                            ),
630                            build_source_map_entry_copy(2, 4),
631                            build_source_map_entry_copy(3, 5),
632                        ],
633                    ),
634                ),
635                build_source_map_entry_copy(2, 7),
636                build_source_map_entry_copy(3, 8),
637            ];
638            Self {
639                program: QUIL,
640                filter: Box::new(|_| true),
641                expected: Ok(EXPECTED_QUIL),
642                source_map_entry_builders,
643            }
644        }
645
646        /// Test that a sequence gate definition works even if one of the qubit parameters
647        /// is not used in the expansion.
648        ///
649        /// This is not expressly forbidden by the Quil specification, so we include
650        /// this test to document the behavior.
651        fn unused_instruction() -> Self {
652            const QUIL: &str = r"
653DEFGATE seq1(%param1) a b AS SEQUENCE:
654    RZ(%param1) a
655    RX(pi/2) a
656    RZ(%param1) a
657
658seq1(pi) 0 1
659";
660            const EXPECTED_QUIL: &str = r"
661RZ(pi) 0
662RX(pi/2) 0
663RZ(pi) 0
664";
665            let source_map_entry_builders = vec![build_source_map_entry(
666                0,
667                DefGateSequenceExpansionBuilder::new(
668                    "seq1",
669                    &["param1"],
670                    &["a", "b"],
671                    0..3,
672                    vec![
673                        build_source_map_entry_copy(0, 0),
674                        build_source_map_entry_copy(1, 1),
675                        build_source_map_entry_copy(2, 2),
676                    ],
677                ),
678            )];
679            Self {
680                program: QUIL,
681                filter: Box::new(|_| true),
682                expected: Ok(EXPECTED_QUIL),
683                source_map_entry_builders,
684            }
685        }
686
687        /// Test sequence gate definition expansion within a program, where one or more
688        /// sequence gate definitions are not included by the filter.
689        fn filtered_sequence() -> Self {
690            const QUIL: &str = r"
691DEFGATE seq1(%param1) a AS SEQUENCE:
692    RZ(%param1) a
693    RX(pi/2) a
694    RZ(%param1) a
695
696DEFGATE seq2(%param1) a AS SEQUENCE:
697    X a
698
699seq1(pi) 0
700seq2(pi/2) 0
701";
702            const EXPECTED_QUIL: &str = r"
703RZ(pi) 0
704RX(pi/2) 0
705RZ(pi) 0
706seq2(pi/2) 0
707";
708            let source_map_entry_builders = vec![
709                build_source_map_entry(
710                    0,
711                    DefGateSequenceExpansionBuilder::new(
712                        "seq1",
713                        &["param1"],
714                        &["a"],
715                        0..3,
716                        vec![
717                            build_source_map_entry_copy(0, 0),
718                            build_source_map_entry_copy(1, 1),
719                            build_source_map_entry_copy(2, 2),
720                        ],
721                    ),
722                ),
723                build_source_map_entry_copy(1, 3),
724            ];
725            Self {
726                program: QUIL,
727                filter: Box::new(|k| k == "seq1"),
728                expected: Ok(EXPECTED_QUIL),
729                source_map_entry_builders,
730            }
731        }
732
733        fn error_parameter_count() -> Self {
734            const QUIL: &str = r"
735DEFGATE seq1(%param1) a AS SEQUENCE:
736    RZ(%param1) a
737seq1() 0
738";
739            let expected = Err(DefGateSequenceExpansionError::ParameterCount {
740                expected: 1,
741                found: 0,
742            });
743            Self {
744                program: QUIL,
745                filter: Box::new(|_| true),
746                expected,
747                source_map_entry_builders: vec![],
748            }
749        }
750
751        fn error_cyclic_sequence_gate_definition() -> Self {
752            const QUIL: &str = r"
753DEFGATE seq1(%param1) a AS SEQUENCE:
754    seq2(%param1) a
755
756DEFGATE seq2(%param1) a AS SEQUENCE:
757    seq3(%param1) a
758
759DEFGATE seq3(%param1) a AS SEQUENCE:
760    seq1(%param1) a
761
762seq1(pi) 0
763";
764            let expected = Err(DefGateSequenceExpansionError::CyclicSequenceGateDefinition(
765                vec!["seq1".to_string(), "seq2".to_string(), "seq3".to_string()],
766            ));
767            Self {
768                program: QUIL,
769                filter: Box::new(|_| true),
770                expected,
771                source_map_entry_builders: vec![],
772            }
773        }
774
775        fn error_qubit_count() -> Self {
776            const QUIL: &str = r"
777DEFGATE seq1(%param1) a AS SEQUENCE:
778    RZ(%param1) a
779
780seq1(pi/2) 0 1
781";
782            let expected = Err(DefGateSequenceExpansionError::QubitCount {
783                expected: 1,
784                found: 2,
785            });
786            Self {
787                program: QUIL,
788                filter: Box::new(|_| true),
789                expected,
790                source_map_entry_builders: vec![],
791            }
792        }
793
794        fn error_gate_qubit_argument() -> Self {
795            const QUIL: &str = r"
796DEFGATE seq1(%param1) a AS SEQUENCE:
797    RZ(%param1) a
798
799seq1(pi/2) %q1
800";
801            let expected = Err(DefGateSequenceExpansionError::NonFixedQubitArgument(
802                crate::instruction::Qubit::Variable("q1".to_string()),
803            ));
804            Self {
805                program: QUIL,
806                filter: Box::new(|_| true),
807                expected,
808                source_map_entry_builders: vec![],
809            }
810        }
811
812        fn error_gate_modifiers_unsupported() -> Self {
813            const QUIL: &str = r"
814DEFGATE seq1(%param1) a AS SEQUENCE:
815    RZ(%param1) a
816
817
818DAGGER seq1(pi/2) 0
819";
820            let expected = Err(DefGateSequenceExpansionError::GateModifiersUnsupported(
821                vec![crate::instruction::GateModifier::Dagger],
822            ));
823            Self {
824                program: QUIL,
825                filter: Box::new(|_| true),
826                expected,
827                source_map_entry_builders: vec![],
828            }
829        }
830    }
831
832    #[rstest]
833    #[case::simple_1q_expansions(DefGateSequenceExpansionTestCase::simple_1q_expansions())]
834    #[case::triple_recursize(DefGateSequenceExpansionTestCase::triple_recursize())]
835    #[case::unexpanded_instructions(DefGateSequenceExpansionTestCase::unexpanded_instructions())]
836    #[case::unused_instruction(DefGateSequenceExpansionTestCase::unused_instruction())]
837    #[case::filtered_sequence(DefGateSequenceExpansionTestCase::filtered_sequence())]
838    #[case::error_qubit_count(DefGateSequenceExpansionTestCase::error_qubit_count())]
839    #[case::error_gate_qubit_argument(DefGateSequenceExpansionTestCase::error_parameter_count())]
840    #[case::error_gate_qubit_argument(DefGateSequenceExpansionTestCase::error_gate_qubit_argument())]
841    #[case::error_gate_modifiers_unsupported(
842        DefGateSequenceExpansionTestCase::error_gate_modifiers_unsupported()
843    )]
844    #[case::error_cyclic_sequence_gate_definition(
845        DefGateSequenceExpansionTestCase::error_cyclic_sequence_gate_definition()
846    )]
847    fn test_defgate_sequence_expansion(#[case] test_case: DefGateSequenceExpansionTestCase) {
848        let program =
849            crate::Program::from_str(test_case.program).expect("must be a valid Quil program");
850        let program_expansion = ProgramDefGateSequenceExpander {
851            gate_definitions: &program.gate_definitions,
852            filter: &test_case.filter,
853        };
854        let result = program_expansion.expand_with_source_map(&program.instructions);
855
856        match (&test_case.expected, result) {
857            (Ok(expected), Ok(result)) => {
858                let expected_program =
859                    Program::from_str(expected).expect("expected program must be valid Quil");
860                let mut actual_program = Program::new();
861                actual_program.add_instructions(result.instructions);
862
863                pretty_assertions::assert_eq!(expected_program, actual_program);
864                pretty_assertions::assert_eq!(test_case.to_source_map(), result.source_map);
865
866                let actual_program_without_source_map = Program::from_instructions(
867                    program_expansion
868                        .expand(&program.instructions)
869                        .expect("expansion without source map should succeed"),
870                );
871                pretty_assertions::assert_eq!(expected_program, actual_program_without_source_map);
872            }
873            (Ok(expected), Err(e)) => {
874                panic!("Expected instructions:\n\n{expected:?}\n\ngot error:\n\n{e:?}");
875            }
876            (Err(expected), Ok(result)) => {
877                panic!(
878                    "Expected error:\n\n{expected:?}\n\ngot:\n\n{:?}",
879                    result.instructions
880                );
881            }
882            (Err(expected), Err(found)) => {
883                pretty_assertions::assert_eq!(*expected, found);
884            }
885        }
886    }
887
888    struct GateSignatureBuilder {
889        gate_name: String,
890        gate_parameters: Vec<String>,
891        gate_qubits: Vec<String>,
892    }
893
894    impl GateSignatureBuilder {
895        fn new(
896            gate_name: &'static str,
897            gate_parameters: &'static [&'static str],
898            gate_qubits: &'static [&'static str],
899        ) -> Self {
900            Self {
901                gate_name: gate_name.to_string(),
902                gate_parameters: gate_parameters.iter().map(|&s| s.to_string()).collect(),
903                gate_qubits: gate_qubits.iter().map(|&s| s.to_string()).collect(),
904            }
905        }
906
907        fn build(&self) -> GateSignature<'_> {
908            GateSignature::try_new(
909                &self.gate_name,
910                &self.gate_parameters,
911                &self.gate_qubits,
912                crate::instruction::GateType::Sequence,
913            )
914            .expect("must be a valid gate signature")
915        }
916    }
917
918    struct DefGateSequenceExpansionBuilder {
919        signature: GateSignatureBuilder,
920        range: Range<usize>,
921        nested_expansions:
922            Vec<SourceMapEntry<InstructionIndex, ExpansionResult<DefGateSequenceExpansionBuilder>>>,
923    }
924
925    impl DefGateSequenceExpansionBuilder {
926        fn new(
927            gate_name: &'static str,
928            gate_parameters: &'static [&'static str],
929            gate_qubits: &'static [&'static str],
930            range: Range<usize>,
931            entries: Vec<
932                SourceMapEntry<InstructionIndex, ExpansionResult<DefGateSequenceExpansionBuilder>>,
933            >,
934        ) -> ExpansionResult<Self> {
935            ExpansionResult::Rewritten(Self {
936                signature: GateSignatureBuilder::new(gate_name, gate_parameters, gate_qubits),
937                range,
938                nested_expansions: entries,
939            })
940        }
941
942        fn build(&self) -> ExpansionResult<DefGateSequenceExpansion<'_>> {
943            let entries: Vec<_> = self
944                .nested_expansions
945                .iter()
946                .map(|entry| SourceMapEntry {
947                    source_location: entry.source_location,
948                    target_location: match entry.target_location() {
949                        ExpansionResult::Rewritten(expansion) => expansion.build(),
950                        ExpansionResult::Unmodified(index) => ExpansionResult::Unmodified(*index),
951                    },
952                })
953                .collect();
954            ExpansionResult::Rewritten(DefGateSequenceExpansion {
955                source_signature: self.signature.build(),
956                range: InstructionIndex(self.range.start)..InstructionIndex(self.range.end),
957                nested_expansions: SourceMap { entries },
958            })
959        }
960    }
961
962    fn build_source_map_entry(
963        source_location: usize,
964        target_location: ExpansionResult<DefGateSequenceExpansionBuilder>,
965    ) -> SourceMapEntry<InstructionIndex, ExpansionResult<DefGateSequenceExpansionBuilder>> {
966        SourceMapEntry {
967            source_location: InstructionIndex(source_location),
968            target_location,
969        }
970    }
971
972    fn build_source_map_entry_copy(
973        source_location: usize,
974        target_location: usize,
975    ) -> SourceMapEntry<InstructionIndex, ExpansionResult<DefGateSequenceExpansionBuilder>> {
976        build_source_map_entry(
977            source_location,
978            ExpansionResult::Unmodified(InstructionIndex(target_location)),
979        )
980    }
981}