ryo-source 0.1.0

//! Macro token parsing utilities for walking struct literals inside macros.
//!
//! This module provides utilities to parse and modify expressions within
//! known macro invocations like `vec![]` and `hashmap!{}`.
//!
//! # Supported Macros
//!
//! - `vec![]` - comma-separated expressions
//! - `hashmap!{}` / `btreemap!{}` - key => value pairs
//!
//! # Example
//!
//! ```ignore
//! use ryo_source::pure::macro_utils;
//!
//! // Parse expressions from vec![]
//! if let Some(exprs) = macro_utils::try_extract_exprs("vec", tokens) {
//!     for expr in &exprs {
//!         // Walk each expression...
//!     }
//! }
//! ```

use super::ast::PureExpr;
use super::convert::ToPure;
use proc_macro2::{TokenStream, TokenTree};
use quote::ToTokens;

/// Known collection macros with comma-separated expressions
const COLLECTION_MACROS: &[&str] = &["vec", "array_vec"];

/// Known map macros with key => value syntax
const MAP_MACROS: &[&str] = &["hashmap", "btreemap", "indexmap"];

/// Try to extract expressions from macro tokens.
///
/// Returns `Some(exprs)` if the macro is a known collection type and
/// parsing succeeds, `None` otherwise.
pub fn try_extract_exprs(name: &str, tokens: &str) -> Option<Vec<PureExpr>> {
    // Normalize macro name (strip path prefix like `std::vec`)
    let macro_name = name.rsplit("::").next().unwrap_or(name);

    if COLLECTION_MACROS.contains(&macro_name) {
        parse_comma_separated(tokens)
    } else if MAP_MACROS.contains(&macro_name) {
        parse_map_values(tokens)
    } else {
        None
    }
}

/// Convert PureExprs back to macro token string.
///
/// Used after modifying expressions to regenerate the macro body.
pub fn exprs_to_tokens(exprs: &[PureExpr]) -> Result<String, super::to_syn::ToSynError> {
    use super::to_syn::ToSyn;

    let parts = exprs
        .iter()
        .map(|e| Ok(e.to_syn()?.to_token_stream().to_string()))
        .collect::<Result<Vec<_>, super::to_syn::ToSynError>>()?;
    Ok(parts.join(", "))
}

/// Convert map entries (key, value pairs) back to macro token string.
///
/// Generates `key => value, ...` format for hashmap!/btreemap! macros.
pub fn map_entries_to_tokens(
    entries: &[(PureExpr, PureExpr)],
) -> Result<String, super::to_syn::ToSynError> {
    use super::to_syn::ToSyn;

    let parts = entries
        .iter()
        .map(|(k, v)| {
            Ok(format!(
                "{} => {}",
                k.to_syn()?.to_token_stream(),
                v.to_syn()?.to_token_stream()
            ))
        })
        .collect::<Result<Vec<_>, super::to_syn::ToSynError>>()?;
    Ok(parts.join(", "))
}

/// Parse comma-separated expressions from token string.
///
/// Handles nested structures (braces, brackets, parens) correctly.
fn parse_comma_separated(tokens: &str) -> Option<Vec<PureExpr>> {
    let ts: TokenStream = tokens.parse().ok()?;
    let chunks = split_by_comma(ts);

    let mut exprs = Vec::new();
    for chunk in chunks {
        if chunk.is_empty() {
            continue;
        }
        match syn::parse2::<syn::Expr>(chunk) {
            Ok(expr) => exprs.push(expr.to_pure()),
            Err(_) => {
                // If any chunk fails to parse, give up on this macro
                return None;
            }
        }
    }

    Some(exprs)
}

/// Parse map macro (hashmap!/btreemap!) values.
///
/// Format: `key => value, key2 => value2, ...`
/// Returns only the values (which may contain struct literals).
fn parse_map_values(tokens: &str) -> Option<Vec<PureExpr>> {
    let ts: TokenStream = tokens.parse().ok()?;
    let entries = split_map_entries(ts)?;

    let mut exprs = Vec::new();
    for (_key, value) in entries {
        match syn::parse2::<syn::Expr>(value) {
            Ok(expr) => exprs.push(expr.to_pure()),
            Err(_) => return None,
        }
    }

    Some(exprs)
}

/// Parse map macro and return key-value pairs for reconstruction.
pub fn try_extract_map_entries(name: &str, tokens: &str) -> Option<Vec<(PureExpr, PureExpr)>> {
    let macro_name = name.rsplit("::").next().unwrap_or(name);

    if !MAP_MACROS.contains(&macro_name) {
        return None;
    }

    let ts: TokenStream = tokens.parse().ok()?;
    let entries = split_map_entries(ts)?;

    let mut result = Vec::new();
    for (key, value) in entries {
        let key_expr = syn::parse2::<syn::Expr>(key).ok()?;
        let value_expr = syn::parse2::<syn::Expr>(value).ok()?;
        result.push((key_expr.to_pure(), value_expr.to_pure()));
    }

    Some(result)
}

/// Split TokenStream by comma, respecting nesting.
fn split_by_comma(ts: TokenStream) -> Vec<TokenStream> {
    let mut chunks = Vec::new();
    let mut current = Vec::new();

    for tt in ts {
        // Check if this is a comma at top level
        let is_comma = matches!(&tt, TokenTree::Punct(p) if p.as_char() == ',');

        if is_comma {
            if !current.is_empty() {
                chunks.push(current.drain(..).collect());
            }
        } else {
            // Groups and other tokens are kept intact
            current.push(tt);
        }
    }

    // Don't forget the last chunk
    if !current.is_empty() {
        chunks.push(current.into_iter().collect());
    }

    chunks
}

/// Split map entries by comma, then each entry by `=>`.
fn split_map_entries(ts: TokenStream) -> Option<Vec<(TokenStream, TokenStream)>> {
    let chunks = split_by_comma(ts);
    let mut entries = Vec::new();

    for chunk in chunks {
        let (key, value) = split_by_fat_arrow(chunk)?;
        entries.push((key, value));
    }

    Some(entries)
}

/// Split a single entry by `=>` (fat arrow).
fn split_by_fat_arrow(ts: TokenStream) -> Option<(TokenStream, TokenStream)> {
    let tokens: Vec<TokenTree> = ts.into_iter().collect();

    // Find `=>` pattern (two consecutive punctuation: '=' and '>')
    let mut split_idx = None;
    let mut i = 0;
    while i < tokens.len().saturating_sub(1) {
        if let (TokenTree::Punct(p1), TokenTree::Punct(p2)) = (&tokens[i], &tokens[i + 1]) {
            if p1.as_char() == '=' && p2.as_char() == '>' {
                split_idx = Some(i);
                break;
            }
        }
        i += 1;
    }

    let idx = split_idx?;
    let key: TokenStream = tokens[..idx].iter().cloned().collect();
    let value: TokenStream = tokens[idx + 2..].iter().cloned().collect();

    Some((key, value))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_vec_simple() {
        let tokens = "1, 2, 3";
        let exprs = try_extract_exprs("vec", tokens).unwrap();
        assert_eq!(exprs.len(), 3);
    }

    #[test]
    fn test_parse_vec_with_structs() {
        let tokens =
            r#"Config { host: "a".into(), port: 80 }, Config { host: "b".into(), port: 81 }"#;
        let exprs = try_extract_exprs("vec", tokens).unwrap();
        assert_eq!(exprs.len(), 2);

        // Verify they are struct expressions
        for expr in &exprs {
            assert!(matches!(expr, PureExpr::Struct { path, .. } if path == "Config"));
        }
    }

    #[test]
    fn test_parse_hashmap_values() {
        let tokens = r#""key1" => Config { port: 80 }, "key2" => Config { port: 81 }"#;
        let exprs = try_extract_exprs("hashmap", tokens).unwrap();
        assert_eq!(exprs.len(), 2);
    }

    #[test]
    fn test_extract_map_entries() {
        let tokens = r#""key1" => Config { port: 80 }, "key2" => Config { port: 81 }"#;
        let entries = try_extract_map_entries("hashmap", tokens).unwrap();
        assert_eq!(entries.len(), 2);

        // Check keys
        assert!(matches!(&entries[0].0, PureExpr::Lit(s) if s.contains("key1")));
        assert!(matches!(&entries[1].0, PureExpr::Lit(s) if s.contains("key2")));
    }

    #[test]
    fn test_unknown_macro_returns_none() {
        let tokens = "1, 2, 3";
        assert!(try_extract_exprs("custom_macro", tokens).is_none());
    }

    #[test]
    fn test_exprs_to_tokens() {
        let exprs = vec![
            PureExpr::Lit("1".to_string()),
            PureExpr::Lit("2".to_string()),
        ];
        let result = exprs_to_tokens(&exprs).unwrap();
        assert!(result.contains("1"));
        assert!(result.contains("2"));
    }
}