qudit_macros/

lib.rs

use proc_macro::TokenStream;
use proc_macro2::{Delimiter, Group, Punct, Spacing, Span, TokenStream as TokenStream2, TokenTree};
use quote::quote;
use syn::{Error, Result};

#[derive(Debug, Clone)]
enum TensorTokens {
    OpenBracket,
    ClosedBracket,
    Comma,
    Number(Vec<TokenTree>),
}

#[derive(Debug, Clone)]
enum RecursiveTensor {
    Scalar(Vec<TokenTree>),
    SubTensor(Vec<RecursiveTensor>),
}

/// Replaces `j` with `c32::new(0.0, 1.0)` in the input token stream.
/// Also makes implicit multiplication explicit. (e.g. `4j` becomes `4.0 * j`)
///
/// # Arguments
///
/// * `input` - A tokenstream containing the input tokens.
///
/// # Returns
///
/// * A tokenstream with the processed tokens.
///
/// # Panics
///
/// * If a literal with suffix or prefix `j` is not a valid number.
///
fn j_processing32(input: TokenStream2) -> TokenStream2 {
    let tokens: Vec<TokenTree> = input.into_iter().collect();
    let mut new_stream = Vec::<TokenTree>::new();
    let mut stream_accumulator: Vec<TokenTree>;

    let mut token: &TokenTree;
    let mut lit_str: String;
    let mut pass: bool;
    for index in 0..tokens.len() {
        stream_accumulator = Vec::<TokenTree>::new();
        pass = false;
        token = &tokens[index];

        // Replaces `j` with c64::new(0.0, 1.0).
        match &token {
            TokenTree::Literal(literal) => {
                lit_str = literal.to_string();
                if let Some(num_part) = lit_str.strip_suffix('j') {
                    if let Ok(number_val) = num_part.parse::<f32>() {
                        stream_accumulator.extend(quote! {#number_val * c32::new(0.0, 1.0)});
                        pass = true;
                    } else {
                        panic!("Not a valid number")
                    }
                } else if let Some(num_part) = lit_str.strip_prefix('j') {
                    if let Ok(number_val) = num_part.parse::<f32>() {
                        stream_accumulator.extend(quote! {#number_val * c32::new(0.0, 1.0)});
                        pass = true;
                    } else {
                        panic!("Not a valid number")
                    }
                }
            }

            TokenTree::Ident(identifier) => {
                if identifier.to_string().as_str() == "j" {
                    stream_accumulator.extend(quote! {c32::new(0.0, 1.0)});
                    pass = true;
                }
            }

            TokenTree::Group(group) => {
                let processed_inner_stream = j_processing32(group.stream());
                let new_group = Group::new(group.delimiter(), processed_inner_stream);
                new_stream.push(TokenTree::Group(new_group));
                continue;
            }

            _ => (),
        }

        if !pass {
            new_stream.push(token.clone());
            continue;
        }

        // Makes implicit multiplication explicit
        pass = false;
        if index > 0 {
            if let TokenTree::Punct(punct) = &tokens[index - 1] {
                match punct.as_char() {
                    ']' | ',' | '+' | '-' | '*' | '/' => pass = true,
                    _ => (),
                }
            }
            if !pass {
                new_stream.push(TokenTree::Punct(Punct::new('*', Spacing::Alone)));
            }
        }

        new_stream.extend(stream_accumulator);

        pass = false;
        if index < tokens.len() - 1 {
            if let TokenTree::Punct(punct) = &tokens[index + 1] {
                match punct.as_char() {
                    ']' | ',' | '+' | '-' | '*' | '/' => pass = true,
                    _ => (),
                }
            }
            if !pass {
                new_stream.push(TokenTree::Punct(Punct::new('*', Spacing::Alone)));
            }
        }
    }
    TokenStream2::from_iter(new_stream)
}

/// Replaces `j` with `c64::new(0.0, 1.0)` in the input token stream.
/// Also makes implicit multiplication explicit. (e.g. `4j` becomes `4.0 * j`)
///
/// # Arguments
///
/// * `input` - A tokenstream containing the input tokens.
///
/// # Returns
///
/// * A tokenstream with the processed tokens.
///
/// # Panics
///
/// * If a literal with suffix or prefix `j` is not a valid number.
///
fn j_processing64(input: TokenStream2) -> TokenStream2 {
    let tokens: Vec<TokenTree> = input.into_iter().collect();
    let mut new_stream = Vec::<TokenTree>::new();
    let mut stream_accumulator: Vec<TokenTree>;

    let mut token: &TokenTree;
    let mut lit_str: String;
    let mut pass: bool;
    for index in 0..tokens.len() {
        stream_accumulator = Vec::<TokenTree>::new();
        pass = false;
        token = &tokens[index];

        // Replaces `j` with c64::new(0.0, 1.0).
        match &token {
            TokenTree::Literal(literal) => {
                lit_str = literal.to_string();
                if let Some(num_part) = lit_str.strip_suffix('j') {
                    if let Ok(number_val) = num_part.parse::<f64>() {
                        stream_accumulator.extend(quote! {#number_val * c64::new(0.0, 1.0)});
                        pass = true;
                    } else {
                        panic!("Not a valid number")
                    }
                } else if let Some(num_part) = lit_str.strip_prefix('j') {
                    if let Ok(number_val) = num_part.parse::<f64>() {
                        stream_accumulator.extend(quote! {#number_val * c64::new(0.0, 1.0)});
                        pass = true;
                    } else {
                        panic!("Not a valid number")
                    }
                }
            }

            TokenTree::Ident(identifier) => {
                if identifier.to_string().as_str() == "j" {
                    stream_accumulator.extend(quote! {c64::new(0.0, 1.0)});
                    pass = true;
                }
            }

            TokenTree::Group(group) => {
                let processed_inner_stream = j_processing64(group.stream());
                let new_group = Group::new(group.delimiter(), processed_inner_stream);
                new_stream.push(TokenTree::Group(new_group));
                continue;
            }

            _ => (),
        }

        if !pass {
            new_stream.push(token.clone());
            continue;
        }

        // Makes implicit multiplication explicit
        pass = false;
        if index > 0 {
            if let TokenTree::Punct(punct) = &tokens[index - 1] {
                match punct.as_char() {
                    ']' | ',' | '+' | '-' | '*' | '/' => pass = true,
                    _ => (),
                }
            }
            if !pass {
                new_stream.push(TokenTree::Punct(Punct::new('*', Spacing::Alone)));
            }
        }

        new_stream.extend(stream_accumulator);

        pass = false;
        if index < tokens.len() - 1 {
            if let TokenTree::Punct(punct) = &tokens[index + 1] {
                match punct.as_char() {
                    ']' | ',' | '+' | '-' | '*' | '/' => pass = true,
                    _ => (),
                }
            }
            if !pass {
                new_stream.push(TokenTree::Punct(Punct::new('*', Spacing::Alone)));
            }
        }
    }
    TokenStream2::from_iter(new_stream)
}

/// Categorizes the tokens in the input token stream to aid in parsing.
///
/// # Arguments
///
/// * `token_stream` - A tokenstream containing the input tokens.
///
/// # Returns
///
/// * A tokenstream with the processed tokens.
///
fn tensor_lexer(token_stream: TokenStream2) -> Result<Vec<TensorTokens>> {
    let mut processed_tokens = Vec::new();
    let token_iterator = token_stream.into_iter();

    let mut number_token_accumulator = Vec::new();

    for token in token_iterator {
        match &token {
            TokenTree::Literal(_literal) => {
                number_token_accumulator.push(token);
            }

            TokenTree::Ident(_identifier) => {
                number_token_accumulator.push(token);
            }

            TokenTree::Punct(punctuation) => match punctuation.as_char() {
                ',' => {
                    if !number_token_accumulator.is_empty() {
                        processed_tokens.push(TensorTokens::Number(number_token_accumulator));
                        number_token_accumulator = Vec::new();
                    }

                    processed_tokens.push(TensorTokens::Comma)
                }
                _ => number_token_accumulator.push(token),
            },

            TokenTree::Group(group) => match group.delimiter() {
                Delimiter::Bracket => {
                    if !number_token_accumulator.is_empty() {
                        processed_tokens.push(TensorTokens::Number(number_token_accumulator));
                        number_token_accumulator = Vec::new();
                    }

                    processed_tokens.push(TensorTokens::OpenBracket);
                    processed_tokens.extend(tensor_lexer(group.stream())?);
                    processed_tokens.push(TensorTokens::ClosedBracket);
                }
                _ => number_token_accumulator.push(token),
            },
        }
    }
    if !number_token_accumulator.is_empty() {
        processed_tokens.push(TensorTokens::Number(number_token_accumulator));
    }
    Ok(processed_tokens)
}

/// Organizes a series of custom tokens into a recursive tensor structure.
///
/// # Arguments
///
/// * `tokens` - A slice of `TensorTokens`, expected from `tensor_lexer`.
///
/// # Returns
///
/// * A recursive tensor storing the user's tokens.
/// * The number of tokens consumed from the input slice.
///
/// # Panics
///
/// * If there is a missing closing bracket.
/// * If the tensor does not start with an opening bracket or is not a scalar.
///
fn tensor_parser(tokens: &[TensorTokens]) -> Result<(RecursiveTensor, usize)> {
    let mut index = 0;

    // A tensor starts with [ or is a scalar.
    if let Some(TensorTokens::OpenBracket) = tokens.get(index) {
        index += 1;

        // Each recursive step is adding one dimension.
        let mut children = Vec::new();
        loop {
            if index > tokens.len() {
                return Err(Error::new(Span::call_site(), "Missing closing bracket"));
            }

            if let Some(TensorTokens::ClosedBracket) = tokens.get(index) {
                index += 1;
                return Ok((RecursiveTensor::SubTensor(children), index));
            }

            let (child, delta_index) = tensor_parser(&tokens[index..])?;
            children.push(child);
            index += delta_index;

            if let Some(TensorTokens::Comma) = tokens.get(index) {
                index += 1;
            }
        }
    } else if let Some(TensorTokens::Number(token_tree_vec)) = tokens.get(index) {
        Ok((RecursiveTensor::Scalar(token_tree_vec.clone()), 1))
    } else {
        Err(Error::new(Span::call_site(), "Not a valid tensor"))
    }
}

/// Flattens the recursive tensor structure into a single vector of tokens and calculates its shape.
///
/// # Arguments
///
/// * `input` - A reference to the recursive tensor structure.
///
/// # Returns
///
/// * A vector containing all elements of the input tensor.
/// * The shape of the input tensor.
///
fn flatten_tensor_data(input: &RecursiveTensor) -> (Vec<TokenStream2>, Vec<usize>) {
    match input {
        RecursiveTensor::Scalar(token_vec) => {
            let stream = TokenStream2::from_iter(token_vec.clone());
            (vec![stream.clone()], vec![])
        }

        RecursiveTensor::SubTensor(subtensors) => {
            // Flattened data calculation
            let mut flat_data = Vec::new();
            for subtensor in subtensors {
                let (mut subtensor_data, _) = flatten_tensor_data(subtensor);
                flat_data.append(&mut subtensor_data);
            }

            // Shape calculation
            let (_, sub_shape) = flatten_tensor_data(&subtensors[0]);
            let mut final_shape = vec![subtensors.len()];
            final_shape.extend(sub_shape);

            (flat_data, final_shape)
        }
    }
}

/// Creates a 64-bit complex tensor from nested brackets. Complex numbers
/// can be created using `j`. (e.g. `4j`, `my_function()j`, or `some_variable * j`)
///
/// # Arguments
///
/// * `input` - The user's desired tensor written in simplified language.
///
/// # Returns
///
/// * A 64-bit complex tensor implementing the user's data.
///
/// # Panics
///
/// * If there is a missing closing bracket.
/// * If the tensor does not start with an opening bracket or is not a scalar.
/// * If a literal with suffix or prefix `j` is not a valid number.
#[proc_macro]
pub fn complex_tensor64(input: TokenStream) -> TokenStream {
    let input_processed = j_processing64(input.into());

    let tokens = match tensor_lexer(input_processed) {
        Ok(inner_val) => inner_val,
        Err(error) => return error.to_compile_error().into(),
    };

    let (recursive_tensor, _) = match tensor_parser(&tokens) {
        Ok(inner_val) => inner_val,
        Err(error) => return error.to_compile_error().into(),
    };

    let (flat_data, shape) = flatten_tensor_data(&recursive_tensor);

    let quoted_shape = quote! {[#(#shape),*]};

    let d = shape.len();
    quote! {{
            let data_vec: Vec<c64> = vec![#(#flat_data),*];
            Tensor::<c64, #d>::from_slice(&data_vec, #quoted_shape)
    }}
    .into()
}

/// Creates a 32-bit complex tensor from nested brackets. Complex numbers
/// can be created using `j`. (e.g. `4j`, `my_function()j`, or `some_variable * j`)
///
/// # Arguments
///
/// * `input` - The user's desired tensor written in simplified language.
///
/// # Returns
///
/// * A 32-bit complex tensor implementing the user's data.
///
/// # Panics
///
/// * If there is a missing closing bracket.
/// * If the tensor does not start with an opening bracket or is not a scalar.
/// * If a literal with suffix or prefix `j` is not a valid number.
#[proc_macro]
pub fn complex_tensor32(input: TokenStream) -> TokenStream {
    let input_processed = j_processing32(input.into());

    let tokens = match tensor_lexer(input_processed) {
        Ok(inner_val) => inner_val,
        Err(error) => return error.to_compile_error().into(),
    };

    let (recursive_tensor, _) = match tensor_parser(&tokens) {
        Ok(inner_val) => inner_val,
        Err(error) => return error.to_compile_error().into(),
    };

    let (flat_data, shape) = flatten_tensor_data(&recursive_tensor);

    let quoted_shape = quote! {[#(#shape),*]};

    let d = shape.len();
    quote! {{
            let data_vec: Vec<c32> = vec![#(#flat_data),*];
            Tensor::<c32, #d>::from_slice(&data_vec, #quoted_shape)
    }}
    .into()
}