candle-einops-macros 0.2.0

use std::collections::hash_map::DefaultHasher;
use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};

use quote::ToTokens;
use syn::parse::discouraged::Speculative;
use syn::{parse::ParseStream, token};

// Custom keywords to represent reduce operations
mod kw {
    syn::custom_keyword!(min);
    syn::custom_keyword!(max);
    syn::custom_keyword!(sum);
    syn::custom_keyword!(mean);
    syn::custom_keyword!(prod);
}

#[derive(Debug, Clone)]
pub enum Decomposition {
    // New decomposed dimension with no shape provided
    // by the user
    Derived {
        name: String,
        index: Index,
        operation: Option<Operation>,
        shape_calc: proc_macro2::TokenStream,
    },
    Named {
        name: String,
        index: Index,
        operation: Option<Operation>,
        shape: Option<Shape>,
    },
}

#[derive(Debug, Clone)]
pub enum Shape {
    Lit(usize),
    Expr(proc_macro2::TokenStream),
}

#[derive(Debug)]
pub enum Composition {
    Individual(Index),
    // Start and end index of dimesions to be combined
    // into one
    Combined { from: Index, to: Option<Index> },
}

#[derive(Debug, Clone)]
pub enum Index {
    Known(usize),
    // Position of dimension after '..' symbol
    Unknown(usize),
    // Position of the '..' symbol
    Range(usize),
}

impl Ord for Index {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match (self, other) {
            (
                Index::Known(i) | Index::Unknown(i) | Index::Range(i),
                Index::Known(j) | Index::Unknown(j) | Index::Range(j),
            ) => i.cmp(j),
        }
    }
}

impl PartialOrd for Index {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Index {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (
                Index::Known(i) | Index::Unknown(i) | Index::Range(i),
                Index::Known(j) | Index::Unknown(j) | Index::Range(j),
            ) => i == j,
        }
    }
}

impl Eq for Index {}

#[derive(Debug, Clone)]
pub enum Operation {
    Min,
    Max,
    Sum,
    Mean,
    Prod,
}

pub fn parse_decomposition(input: ParseStream) -> syn::Result<(Vec<Decomposition>, bool, usize)> {
    // Length of dimensions inside parenthesis,
    // we need it to account for the dimensions skipped during the iteration
    let mut parenthesized_len = 0;
    let mut minimum_rank = 0;
    let (decomposition, requires_decomposition, _) = (0..)
        // We parse till we reach '->'. The token is consumed after the fold so
        // parsing failures can be propagated instead of panicking here.
        .take_while(|_| !input.peek(syn::Token![->]))
        .try_fold(
            (
                Vec::new(),
                // Boolean to indicate if,
                // - dimensions are decomposed
                // - squeezing is required
                false,
                // Closure that helps construct `Index`,
                // it is updated once we hit '..' to construct `Unknown` indices
                Box::new(Index::Known) as Box<dyn Fn(usize) -> Index>,
            ),
            |(mut decomposition, mut requires_decomposition, mut index_fn), mut i| {
                // We account for dimensions skipped when we parse
                // contents of parenthesized expression
                i += parenthesized_len;
                if input.peek(syn::token::Paren) {
                    let content_expression = parse_left_parenthesized(input, index_fn(i))?;
                    decomposition.extend(content_expression);
                    requires_decomposition = true;
                    minimum_rank += 1;
                } else if peek_reduce_kw(input) {
                    let span = input.span();
                    let identifiers = parse_reduce_fn(input)?;
                    if identifiers.iter().any(|(_, shape, _)| shape.is_some()) {
                        return Err(syn::Error::new(
                            span,
                            "Axis sizes are not allowed in top-level reductions",
                        ));
                    }
                    minimum_rank += identifiers.iter().filter(|(name, ..)| name != "..").count();
                    // We account for dimensions skipped during reduction operations
                    // like `sum(a b c)`, where three dimensions are reduced
                    parenthesized_len += identifiers.len().saturating_sub(1);
                    identifiers.into_iter().enumerate().for_each(
                        |(inner_index, (name, shape, operation))| {
                            // If we encounter '..' inside reduce, we will have to update the closure
                            if name == ".." {
                                index_fn = Box::new(Index::Unknown);
                                decomposition.push(Decomposition::Named {
                                    name,
                                    index: Index::Range(i + inner_index),
                                    shape,
                                    operation: Some(operation),
                                });
                            } else {
                                decomposition.push(Decomposition::Named {
                                    name,
                                    index: index_fn(i + inner_index),
                                    shape,
                                    operation: Some(operation),
                                });
                            }
                        },
                    );
                } else if input.peek(syn::Ident) {
                    let span = input.span();
                    let (name, shape) = parse_identifier(input)?;
                    if shape.is_some() {
                        return Err(syn::Error::new(
                            span,
                            "Axis sizes are only allowed inside a decomposition group on the left",
                        ));
                    }
                    decomposition.push(Decomposition::Named {
                        name,
                        shape,
                        index: index_fn(i),
                        operation: None,
                    });
                    minimum_rank += 1;
                } else if input.peek(syn::LitInt) {
                    let lit_int = input.parse::<syn::LitInt>()?;
                    if lit_int.base10_parse::<usize>()? != 1 {
                        // '1' is the only literal int allowed on the left,
                        // this indicates a squeeze operation
                        return Err(input.error(format!(
                            "Literal Int {} not allowed on the left side",
                            lit_int
                        )));
                    }
                    // We have to reshape to squeeze the 1 sized dimension
                    requires_decomposition = true;
                    minimum_rank += 1;
                } else if input.peek(syn::Token![..]) {
                    input.parse::<syn::Token![..]>()?;
                    decomposition.push(Decomposition::Named {
                        name: "..".to_string(),
                        index: Index::Range(i),
                        shape: None,
                        operation: None,
                    });
                    // We update the closure as we have encountered '..'
                    index_fn = Box::new(Index::Unknown);
                } else {
                    return Err(input
                        .error("Unrecognized charater found in the left side of the expression"));
                }
                Ok((decomposition, requires_decomposition, index_fn))
            },
        )?;

    input.parse::<syn::Token![->]>()?;

    Ok((decomposition, requires_decomposition, minimum_rank))
}

fn parse_left_parenthesized(input: ParseStream, index: Index) -> syn::Result<Vec<Decomposition>> {
    let span = input.span();
    let content;
    syn::parenthesized!(content in input);

    if content.is_empty() {
        return Err(syn::Error::new(
            span,
            "Decomposition groups cannot be empty",
        ));
    }

    let mut content_expression = Vec::new();

    let (derived_name, derived_index, derived_operation, running_mul, shape_expr) = (0..)
        // We continue till we parse everything inside the parenthesis
        .take_while(|_| !content.is_empty())
        .try_fold(
            (None, None, None, 1usize, Vec::new()),
            |(
                mut derived_name,
                mut derived_index,
                mut derived_operation,
                mut running_mul,
                mut shape_expr,
            ),
             i| {
                // Closure to keep a running multiple of the shapes,
                // and updating the list with new dimensions with known shape
                let mut update_values = |name, shape, operation| {
                    if let Some(size) = shape {
                        match size {
                            Shape::Lit(lit_size) => {
                                running_mul =
                                    running_mul.checked_mul(lit_size).ok_or_else(|| {
                                        syn::Error::new(
                                            span,
                                            "Decomposition size product overflows usize",
                                        )
                                    })?;
                            }
                            Shape::Expr(ref expression) => shape_expr.push(expression.clone()),
                        }
                        content_expression.push(Decomposition::Named {
                            name,
                            index: index.clone(),
                            operation,
                            shape: Some(size),
                        });
                    } else {
                        if derived_name.is_some() {
                            // We cannot have more than one dimension with unknown shape
                            return Err(content
                                .error("Shape information required to complete decomposition"));
                        }
                        derived_name = Some(name);
                        derived_index = Some(i);
                        derived_operation = operation;
                    }
                    Ok(())
                };
                if peek_reduce_kw(&content) {
                    parse_reduce_fn(&content)?.into_iter().try_for_each(
                        |(name, shape, operation)| {
                            if name == ".." {
                                return Err(content.error(
                                    "Ignore symbol '..' not allowed inside brackets on the left",
                                ));
                            }
                            update_values(name, shape, Some(operation))
                        },
                    )?;
                } else if content.peek(syn::Ident) {
                    let (name, shape) = parse_identifier(&content)?;
                    update_values(name, shape, None)?
                } else if content.peek(syn::Token![..]) {
                    return Err(
                        content.error("Ignore symbol '..' not allowed inside brackets on the left")
                    );
                } else if content.peek(syn::LitInt) {
                    let lit_int = content.parse::<syn::LitInt>()?;
                    return Err(content.error(format!(
                        "Anonymous integer {} is not allowed inside brackets on the left",
                        lit_int
                    )));
                } else if content.peek(syn::token::Brace) {
                    let (name, shape) = parse_braced_expression(&content)?;
                    update_values(name, Some(shape), None)?;
                } else {
                    return Err(content.error(
                        "Unknown character found inside the brackets of the left expression",
                    ));
                };
                Ok((
                    derived_name,
                    derived_index,
                    derived_operation,
                    running_mul,
                    shape_expr,
                ))
            },
        )?;

    // We add the `Derived` dimension at the end at its index,
    // once we have the running multiple of all the other shapes
    // inside the parenthesis
    if let Some(derived_index) = derived_index {
        let derived_name = derived_name.ok_or_else(|| {
            syn::Error::new(
                content.span(),
                "Internal error while deriving the decomposed axis name",
            )
        })?;
        content_expression.insert(
            derived_index,
            Decomposition::Derived {
                name: derived_name,
                index,
                operation: derived_operation,
                shape_calc: if shape_expr.is_empty() {
                    quote::quote!(::core::option::Option::Some(#running_mul))
                } else {
                    quote::quote!(
                        [#(#shape_expr),*]
                            .into_iter()
                            .try_fold(#running_mul, |product, factor| product.checked_mul(factor))
                    )
                },
            },
        );
    }

    Ok(content_expression)
}

fn parse_braced_expression(content: ParseStream) -> syn::Result<(String, Shape)> {
    let shape_expression;
    syn::braced!(shape_expression in content);

    let span = shape_expression.span();

    let fork = shape_expression.fork();
    let (braced_name, braced_shape) = if let Ok(field) = fork.parse::<syn::ExprField>() {
        shape_expression.advance_to(&fork);
        let mut hasher = DefaultHasher::new();
        field.hash(&mut hasher);
        (
            hasher.finish().to_string(),
            Shape::Expr(field.to_token_stream()),
        )
    } else if shape_expression.peek(syn::Ident) {
        let ident = shape_expression.parse::<syn::Ident>()?;
        (ident.to_string(), Shape::Expr(ident.to_token_stream()))
    } else {
        return Err(syn::Error::new(
            span,
            "Only identifiers and fields are allowed inside '{}'",
        ));
    };

    if !shape_expression.is_empty() {
        return Err(syn::Error::new(
            span,
            "Only one identifier or field is allowed inside '{}'",
        ));
    }

    Ok((braced_name, braced_shape))
}

pub fn parse_reduce(decomposition: &[Decomposition]) -> Vec<(Index, Operation)> {
    // We filter for only the dimensions that have some type of
    // operation associated with them
    decomposition
        .iter()
        .cloned()
        .enumerate()
        .filter_map(|(i, expression)| match expression {
            Decomposition::Named {
                index: Index::Known(_),
                operation: Some(operation),
                ..
            }
            | Decomposition::Derived {
                index: Index::Known(_),
                operation: Some(operation),
                ..
            } => Some((Index::Known(i), operation)),
            Decomposition::Named {
                index: Index::Unknown(_),
                operation: Some(operation),
                ..
            }
            | Decomposition::Derived {
                index: Index::Unknown(_),
                operation: Some(operation),
                ..
            } => Some((Index::Unknown(i), operation)),
            Decomposition::Named {
                index: Index::Range(_),
                operation: Some(operation),
                ..
            } => Some((Index::Range(i), operation)),
            _ => None,
        })
        .collect::<Vec<_>>()
}

#[allow(clippy::type_complexity)]
pub fn parse_composition_permute_repeat(
    input: ParseStream,
    decomposition: &[Decomposition],
) -> syn::Result<(Vec<Composition>, Vec<Index>, Vec<(Index, Shape)>)> {
    // We calculate the span to report errors later
    let input_span = input.span();
    let mut left_names = HashSet::new();
    for expression in decomposition {
        let (name, is_anonymous_literal) = match expression {
            Decomposition::Named {
                name,
                shape: Some(Shape::Lit(size)),
                ..
            } => (name, name == &size.to_string()),
            Decomposition::Named { name, .. } | Decomposition::Derived { name, .. } => {
                (name, false)
            }
        };
        // Literal axes such as `max(2)` are anonymous and may be reused.
        if is_anonymous_literal {
            continue;
        }
        if !left_names.insert(name) {
            let message = if name == ".." {
                "Ellipsis `..` appears more than once on the left".to_string()
            } else {
                format!("Axis `{name}` appears more than once on the left")
            };
            return Err(syn::Error::new(input_span, message));
        }
    }
    // We check if ignored dimensions are reduced
    let is_ignore_reduced = decomposition.iter().any(|expression| {
        matches!(expression, Decomposition::Named {name, operation: Some(_), ..} if name.as_str() == "..")
    });
    // If '..' is reduced, previously unknown indices could change
    // to known
    let unknown_index_fn = |i| {
        if is_ignore_reduced {
            return Index::Known(i);
        }
        Index::Unknown(i)
    };
    // We create a hashmap of identifiers and their positions
    // in the left expression
    let positions = decomposition
        .iter()
        // We ignore all the reduced dimensions
        .filter(|expression| {
            !matches!(
                expression,
                Decomposition::Named {
                    operation: Some(_),
                    ..
                } | Decomposition::Derived {
                    operation: Some(_),
                    ..
                }
            )
        })
        .enumerate()
        .try_fold(HashMap::new(), |mut map, (i, expression)| {
            let (name, index) = match expression {
                Decomposition::Named {
                    name,
                    index: Index::Known(_),
                    ..
                }
                | Decomposition::Derived {
                    name,
                    index: Index::Known(_),
                    ..
                } => (name, Index::Known(i)),
                Decomposition::Named {
                    name,
                    index: Index::Unknown(_),
                    ..
                }
                | Decomposition::Derived {
                    name,
                    index: Index::Unknown(_),
                    ..
                } => (name, unknown_index_fn(i)),
                Decomposition::Named {
                    name,
                    index: Index::Range(_),
                    ..
                } => (name, Index::Range(i)),
                _ => unreachable!(),
            };
            if map.insert(name.clone(), index).is_some() {
                return Err(syn::Error::new(
                    input_span,
                    format!("Axis `{name}` appears more than once on the left"),
                ));
            }
            Ok(map)
        })?;

    let mut consumed = HashSet::new();

    // To keep track of the dimensions skipped inside parenthesis
    let mut parenthesized_len: usize = 0;
    let (composition, permute, repeat, _) = (0..)
        .take_while(|_| !input.is_empty())
        .try_fold::<_, _, syn::Result<(_, _, _, _)>>(
            (
                Vec::new(),
                Vec::new(),
                Vec::new(),
                // Closure to construct `Index`, once we encounter '..',
                // `Known` index becomes `Unknown`
                Box::new(Index::Known) as Box<dyn Fn(usize) -> Index>,
            ),
            |(mut composition, mut permute, mut repeat, mut index_fn), mut i| {
                // We update the iterator's index to account for dimensions skipped
                // inside parenthesis
                i += parenthesized_len;
                if input.peek(token::Paren) {
                    let (combined, combined_permute, combined_repeat, combined_len) =
                        parse_right_parenthesized(
                            input,
                            i,
                            &mut index_fn,
                            &positions,
                            &mut consumed,
                        )?;
                    parenthesized_len += combined_len.saturating_sub(1);
                    permute.extend(combined_permute);
                    repeat.extend(combined_repeat);
                    composition.push(combined);
                } else if input.peek(syn::Ident) {
                    let (name, shape) = parse_identifier(input)?;
                    if let Some(index) = positions.get(&name) {
                        if shape.is_some() {
                            return Err(syn::Error::new(
                                input_span,
                                format!("Axis `{name}` cannot be assigned a size on the right"),
                            ));
                        }
                        permute.push(index.clone());
                    } else {
                        // New identifiers represents repetition
                        let shape = shape.ok_or_else(|| {
                            syn::Error::new(
                                input.span(),
                                format!("New axis `{name}` requires an explicit size"),
                            )
                        })?;
                        repeat.push((index_fn(i), shape));
                    }
                    if !consumed.insert(name.clone()) {
                        return Err(syn::Error::new(
                            input_span,
                            format!("Axis `{name}` appears more than once on the right"),
                        ));
                    }
                    composition.push(Composition::Individual(index_fn(i)))
                } else if input.peek(syn::LitInt) {
                    // Literal ints represent repetition
                    repeat.push((index_fn(i), Shape::Lit(parse_usize(input)?)));
                    composition.push(Composition::Individual(index_fn(i)));
                } else if input.peek(syn::Token![..]) {
                    input.parse::<syn::Token![..]>()?;
                    if !consumed.insert("..".to_string()) {
                        return Err(syn::Error::new(
                            input_span,
                            "Ellipsis `..` appears more than once on the right",
                        ));
                    }
                    composition.push(Composition::Individual(Index::Range(i)));
                    let index = positions.get("..").ok_or_else(|| {
                        syn::Error::new(
                            input.span(),
                            "Ellipsis `..` must appear on both sides of the expression",
                        )
                    })?;
                    permute.push(index.clone());
                    // We update the closure
                    index_fn = Box::new(Index::Unknown);
                } else if input.peek(syn::token::Brace) {
                    let (name, shape) = parse_braced_expression(input)?;
                    if !consumed.insert(name.clone()) {
                        return Err(syn::Error::new(
                            input_span,
                            format!("Axis `{name}` appears more than once on the right"),
                        ));
                    }
                    if let Some(index) = positions.get(&name) {
                        permute.push(index.clone());
                    } else {
                        repeat.push((index_fn(i), shape));
                    }
                    composition.push(Composition::Individual(index_fn(i)));
                } else {
                    return Err(
                        input.error("Unrecognized character on the right side of the expression")
                    );
                }
                Ok((composition, permute, repeat, index_fn))
            },
        )?;

    // We raise an error if the right side of the expression
    // misses a indicator from the left side
    let mut missing = positions
        .keys()
        .filter(|name| !consumed.contains(*name))
        .cloned()
        .collect::<Vec<_>>();
    missing.sort();
    if !missing.is_empty() {
        return Err(syn::Error::new(
            input_span,
            format!("Axes missing on the right: {}", missing.join(", ")),
        ));
    }

    Ok((composition, permute, repeat))
}

#[allow(clippy::type_complexity)]
fn parse_right_parenthesized(
    input: ParseStream,
    start_index: usize,
    index_fn: &mut Box<dyn Fn(usize) -> Index>,
    positions: &HashMap<String, Index>,
    consumed: &mut HashSet<String>,
) -> syn::Result<(Composition, Vec<Index>, Vec<(Index, Shape)>, usize)> {
    let span = input.span();
    let content;
    syn::parenthesized!(content in input);

    if content.is_empty() {
        return Err(syn::Error::new(span, "Composition groups cannot be empty"));
    }

    let mut permute = Vec::new();
    let mut repeat = Vec::new();

    // Closure to parse one entry in the expression and update
    // the relevant lists
    let mut parse_content = |content: ParseStream, index: usize| -> syn::Result<Index> {
        if content.peek(syn::Token![..]) {
            content.parse::<syn::Token![..]>()?;
            if !consumed.insert("..".to_string()) {
                return Err(syn::Error::new(
                    span,
                    "Ellipsis `..` appears more than once on the right",
                ));
            }
            let ignored_index = positions.get("..").ok_or_else(|| {
                syn::Error::new(
                    content.span(),
                    "Ellipsis `..` must appear on both sides of the expression",
                )
            })?;
            permute.push(ignored_index.clone());
            *index_fn = Box::new(Index::Unknown);
            Ok(Index::Range(index))
        } else if content.peek(syn::Ident) {
            let (name, shape) = parse_identifier(content)?;
            if let Some(index) = positions.get(&name) {
                if shape.is_some() {
                    return Err(syn::Error::new(
                        span,
                        format!("Axis `{name}` cannot be assigned a size on the right"),
                    ));
                }
                permute.push(index.clone());
            } else {
                let shape = shape.ok_or_else(|| {
                    syn::Error::new(
                        content.span(),
                        format!("New axis `{name}` requires an explicit size"),
                    )
                })?;
                repeat.push((index_fn(index), shape));
            }
            if !consumed.insert(name.clone()) {
                return Err(syn::Error::new(
                    span,
                    format!("Axis `{name}` appears more than once on the right"),
                ));
            }
            Ok(index_fn(index))
        } else if content.peek(syn::LitInt) {
            repeat.push((index_fn(index), Shape::Lit(parse_usize(content)?)));
            Ok(index_fn(index))
        } else if content.peek(syn::token::Brace) {
            let (name, shape) = parse_braced_expression(content)?;
            if !consumed.insert(name.clone()) {
                return Err(syn::Error::new(
                    span,
                    format!("Axis `{name}` appears more than once on the right"),
                ));
            }
            if let Some(index) = positions.get(&name) {
                permute.push(index.clone());
            } else {
                repeat.push((index_fn(index), shape));
            }
            Ok(index_fn(index))
        } else {
            Err(input.error("Unrecognized character on the right side of the expression"))
        }
    };

    // The starting index of the combined dimensions
    let from = parse_content(&content, start_index)?;

    // The ending index of the combined dimension,
    // we iterate through the entire expression to update relevant lists
    let mut to = None;
    for i in (start_index + 1).. {
        if content.is_empty() {
            break;
        }
        to = Some(parse_content(&content, i)?);
    }

    // We calculate the length of dimesions inside this parenthesis,
    // this helps the main loop keep track of skipped dimensions
    let len = if let Some(
        Index::Known(end_index) | Index::Unknown(end_index) | Index::Range(end_index),
    ) = to
    {
        (end_index - start_index) + 1
    } else {
        0
    };

    Ok((Composition::Combined { from, to }, permute, repeat, len))
}

fn peek_reduce_kw(input: ParseStream) -> bool {
    input.peek(kw::min)
        | input.peek(kw::max)
        | input.peek(kw::sum)
        | input.peek(kw::mean)
        | input.peek(kw::prod)
}

fn parse_reduce_fn(input: ParseStream) -> syn::Result<Vec<(String, Option<Shape>, Operation)>> {
    let span = input.span();
    let operation = if input.peek(kw::min) {
        input.parse::<kw::min>()?;
        Operation::Min
    } else if input.peek(kw::max) {
        input.parse::<kw::max>()?;
        Operation::Max
    } else if input.peek(kw::sum) {
        input.parse::<kw::sum>()?;
        Operation::Sum
    } else if input.peek(kw::mean) {
        input.parse::<kw::mean>()?;
        Operation::Mean
    } else if input.peek(kw::prod) {
        input.parse::<kw::prod>()?;
        Operation::Prod
    } else {
        return Err(input.error("Expected a reduction operation"));
    };

    let content;
    syn::parenthesized!(content in input);

    if content.is_empty() {
        return Err(syn::Error::new(span, "Reduction groups cannot be empty"));
    }

    Ok(content
        // A single operation call can have multiple dimensions,
        .call(parse_identifiers)?
        .into_iter()
        .map(|(name, shape)| (name, shape, operation.clone()))
        .collect())
}

fn parse_identifiers(content: ParseStream) -> syn::Result<Vec<(String, Option<Shape>)>> {
    let mut identifiers = Vec::new();
    while !content.is_empty() {
        if content.peek(syn::Ident) {
            let (name, shape) = content.call(parse_identifier)?;
            identifiers.push((name, shape));
        } else if content.peek(syn::Token![..]) {
            content.parse::<syn::Token![..]>()?;
            identifiers.push(("..".to_string(), None));
        } else if content.peek(syn::LitInt) {
            let lit_int = parse_usize(content)?;
            identifiers.push((lit_int.to_string(), Some(Shape::Lit(lit_int))));
        } else if content.peek(syn::token::Brace) {
            let (name, shape) = parse_braced_expression(content)?;
            identifiers.push((name, Some(shape)));
        } else {
            return Err(content.error("Unknown character introduced in the reduce operation"));
        }
    }
    Ok(identifiers)
}

fn parse_identifier(input: ParseStream) -> syn::Result<(String, Option<Shape>)> {
    let name = input.parse::<syn::Ident>()?.to_string();

    let shape = if input.peek(syn::Token![:]) {
        input.parse::<syn::Token![:]>()?;
        Some(Shape::Lit(parse_usize(input)?))
    } else {
        None
    };

    Ok((name, shape))
}

fn parse_usize(input: ParseStream) -> syn::Result<usize> {
    let len = input.parse::<syn::LitInt>()?;
    len.base10_parse::<usize>()
}