use crate::{extract_fields_with_types, extract_variants_with_types};
use proc_macro2::TokenStream;
use quote::quote;
use syn::{Data, DataEnum, DataStruct, DeriveInput, Ident};
pub fn dispatch_struct_enum(input: DeriveInput) -> TokenStream {
let ident = input.ident;
match input.data {
Data::Struct(ds) => involve_qubits_struct(ds, ident),
Data::Enum(de) => involve_qubits_enum(de, ident),
_ => panic!("InvolveQubits can only be derived on structs and enums"),
}
}
fn involve_qubits_enum(de: DataEnum, ident: Ident) -> TokenStream {
let variants_with_type = extract_variants_with_types(de).into_iter();
let match_quotes = variants_with_type.clone().map(|(vident, _, _)| {
quote! {
&#ident::#vident(ref inner) => {InvolveQubits::involved_qubits(&(*inner))},
}
});
let match_quotes_classical = variants_with_type.map(|(vident, _, _)| {
quote! {
&#ident::#vident(ref inner) => {InvolveQubits::involved_classical(&(*inner))},
}
});
quote! {
#[automatically_derived]
impl InvolveQubits for #ident{
fn involved_qubits(&self) -> InvolvedQubits {
match self{
#(#match_quotes)*
_ => panic!("Unexpectedly cannot match variant")
}
}
fn involved_classical(&self) -> InvolvedClassical {
match self{
#(#match_quotes_classical)*
_ => panic!("Unexpectedly cannot match variant")
}
}
}
}
}
fn involve_qubits_struct(ds: DataStruct, ident: Ident) -> TokenStream {
let fields_with_type = extract_fields_with_types(ds).into_iter();
let mut qubit: bool = false;
let mut control: bool = false;
let mut control_0: bool = false;
let mut control_1: bool = false;
let mut target: bool = false;
let mut qubits: bool = false;
for (fid, type_string, _) in fields_with_type {
match fid.clone().to_string().as_str() {
"qubit" => {
if type_string == Some("usize".to_string()) {
qubit = true
} else {
panic!("Field qubit must have type usize")
}
}
"target" => {
if type_string == Some("usize".to_string()) {
target = true;
} else {
panic!("Field target must have type usize")
}
}
"control" => {
if type_string == Some("usize".to_string()) {
control = true;
} else {
panic!("Field control must have type usize")
}
}
"control_0" => {
if type_string == Some("usize".to_string()) {
control_0 = true;
} else {
panic!("Field control_0 must have type usize")
}
}
"control_1" => {
if type_string == Some("usize".to_string()) {
control_1 = true;
} else {
panic!("Field control_1 must have type usize")
}
}
"qubits" => {
qubits = true;
}
_ => {}
};
}
if qubit {
if control || target || qubits {
panic!("When deriving InvolveQubits, qubit field is not compatible with control, target or qubits fields");
};
quote! {
#[automatically_derived]
impl InvolveQubits for #ident{
fn involved_qubits(&self ) -> InvolvedQubits {
let mut new_hash_set: std::collections::HashSet<usize> = std::collections::HashSet::new();
new_hash_set.insert(self.qubit);
InvolvedQubits::Set(new_hash_set)
}
}
}
} else if control_0 || control_1 {
if control {
panic!("When deriving InvolveQubits for a three-qubit operation, control field must not be present")
}
if !(control_0 && control_1 && target) {
panic!("When deriving InvolveQubits for a three-qubit operation control_0, control_1 and target have to be present");
};
if qubits {
panic!("When deriving InvolveQubits for a three-qubit operation, control_0 and control_1 fields are not compatible with qubits fields");
};
quote! {
#[automatically_derived]
impl InvolveQubits for #ident{
fn involved_qubits(&self ) -> InvolvedQubits {
let mut new_hash_set: std::collections::HashSet<usize> = std::collections::HashSet::new();
new_hash_set.insert(self.control_0);
new_hash_set.insert(self.control_1);
new_hash_set.insert(self.target);
InvolvedQubits::Set(new_hash_set)
}
}
}
} else if target || control {
if !(control && target) {
panic!("When deriving InvolveQubits control and target fields have to both be present");
};
if qubits {
panic!("When deriving InvolveQubits, control and target fields are not compatible with qubits fields");
};
quote! {
#[automatically_derived]
impl InvolveQubits for #ident{
fn involved_qubits(&self ) -> InvolvedQubits {
let mut new_hash_set: std::collections::HashSet<usize> = std::collections::HashSet::new();
new_hash_set.insert(self.control);
new_hash_set.insert(self.target);
InvolvedQubits::Set(new_hash_set)
}
}
}
} else if qubits {
quote! {
#[automatically_derived]
impl InvolveQubits for #ident{
fn involved_qubits(&self ) -> InvolvedQubits {
let mut new_hash_set: std::collections::HashSet<usize> = std::collections::HashSet::new();
for qubit in self.qubits.iter(){
new_hash_set.insert(*qubit);
}
InvolvedQubits::Set(new_hash_set)
}
}
}
} else {
panic!("To derive InvolveQubits qubit or control or target or qubits fields need to be present in struct")
}
}