1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

// Copyright © 2021-2022 HQS Quantum Simulations GmbH. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the
// License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing permissions and
// limitations under the License.

use quote::format_ident;
use std::collections::HashSet;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{Ident, Token, Type, TypePath};

mod product_wrapper;
use product_wrapper::productwrapper;

/// Attribute macro for constructing the pyo3 implementation for mixed indices.
#[proc_macro_attribute]
pub fn product_wrapper(
    metadata: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    productwrapper(metadata, input)
}

mod noiseless_system_wrapper;
use noiseless_system_wrapper::noiselesswrapper;

/// Attribute macro for constructing the pyo3 implementation for mixed indices.
#[proc_macro_attribute]
pub fn noiseless_system_wrapper(
    metadata: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    noiselesswrapper(metadata, input)
}

mod noisy_system_wrapper;
use noisy_system_wrapper::noisywrapper;

/// Attribute macro for constructing the pyo3 implementation for mixed indices.
#[proc_macro_attribute]
pub fn noisy_system_wrapper(
    metadata: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    noisywrapper(metadata, input)
}

// Helper functions
// Struct for parsed derive macro arguments. Used to identify structs belonging to enums
#[derive(Debug)]
struct AttributeMacroArguments(HashSet<String>);

impl AttributeMacroArguments {
    pub fn contains(&self, st: &str) -> bool {
        self.0.contains(st)
    }
    pub fn _ids(&self) -> Vec<Ident> {
        self.0
            .clone()
            .into_iter()
            .map(|s| format_ident!("{}", s))
            .collect()
    }
}

fn strip_python_wrapper_name(ident: &Type) -> (String, proc_macro2::Ident) {
    // get name of the interal struct (not the wrapper)
    let type_path = match ident.clone() {
        Type::Path(TypePath { path: p, .. }) => p,
        _ => panic!("Trait only supports newtype variants with normal types of form path"),
    };
    let type_string = match type_path.get_ident() {
        Some(ident_path) => ident_path.to_string(),
        _ => match type_path.segments.last() {
            Some(segment) => segment.ident.to_string(),
            None => panic!("Can't extract string."),
        },
    };
    // Cut off "Wrapper" at the end of the Impl name
    let struct_name = type_string
        .as_str()
        .strip_suffix("Wrapper")
        .expect("Not conform to Wrapper naming scheme.");
    let struct_ident = format_ident!("{}", struct_name);
    (struct_name.to_string(), struct_ident)
}

impl Parse for AttributeMacroArguments {
    fn parse(input: ParseStream) -> syn::parse::Result<Self> {
        // Parse arguments as comma separated list of idents
        let arguments = Punctuated::<Ident, Token![,]>::parse_terminated(input)?;
        Ok(Self(
            arguments.into_iter().map(|id| id.to_string()).collect(),
        ))
    }
}