adventofcode/macros/src/lib.rs

use std::fs;

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse, Result};
use syn::{
    parse::{Parse, ParseStream},
    spanned::Spanned,
    Error, Ident, Lit, LitInt,
};

const DATA_DIR: &str = "data";
const EXAMPLE_DIR: &str = "examples";
const VAR_NAME: &str = "DATA";

#[derive(Debug, Clone)]
struct IncludeData {
    var: String,
    year: Option<String>,
    day: Option<String>,
}

impl Default for IncludeData {
    fn default() -> Self {
        IncludeData {
            var: VAR_NAME.into(),
            year: None,
            day: None,
        }
    }
}

fn get_text<T: Spanned>(item: T) -> Result<String> {
    let span = item.span();
    match span.source_text() {
        Some(text) => Ok(text),
        None => {
            let start = span.start();
            let end = span.end();
            Err(Error::new(
                span,
                format!(
                    "Failed to get sourcetext for {}:{}-{}:{}",
                    start.line, start.column, end.line, end.column,
                ),
            ))
        }
    }
}

impl Parse for IncludeData {
    fn parse(input: ParseStream) -> Result<Self> {
        let mut data = IncludeData::default();
        if input.peek(Ident) && !input.peek(LitInt) {
            data.var = get_text(input.parse::<Ident>().unwrap())?;
        }
        data.year = Some(get_text(input.parse::<LitInt>()?)?);
        data.day = Some(get_text(input.parse::<Lit>()?)?);

        Ok(data)
    }
}

fn canonicalize(dir: &str, input: IncludeData) -> Option<String> {
    let pathname = format!("{}/{}/{}.txt", dir, input.year.unwrap(), input.day.unwrap());
    match fs::canonicalize(pathname) {
        Ok(canon) => {
            if canon.is_file() {
                match canon.to_str() {
                    Some(c) => Some(c.to_string()),
                    None => None,
                }
            } else {
                None
            }
        }
        Err(_) => None,
    }
}

/// includes Data from Advent of code
/// ```ignore (cannot-doctest-external-file-dependency)
/// include_data!(DATA 2015 01)
/// // or
/// include_data!(2015 01)
///
/// fn main() {
///     print!("{DATA}");
/// }
/// ```
/// The Equivalent is
/// ```ignore (cannot-doctest-external-file-dependency)
/// static DATA: &str = include_str!("../../../data/2015/01.txt");
/// ```
#[proc_macro]
pub fn include_data(data: TokenStream) -> TokenStream {
    let input = match parse::<IncludeData>(data) {
        Ok(include) => include,
        Err(error) => return error.into_compile_error().into(),
    };
    let ident = format_ident!("{}", input.var);
    let mut comment: String;
    if input.day.is_some() && input.year.is_some() {
        comment = "Data from the data dir".into();
        let mut path = canonicalize(DATA_DIR, input.clone());
        if path.is_none() {
            comment = "Data from the examples".into();
            path = canonicalize(EXAMPLE_DIR, input)
        }
        match path {
            Some(p) => {
                return quote! {
                    #[doc = #comment]
                    const #ident: &str = include_str!(#p).trim_ascii();
                }
                .into()
            }
            None => comment = "failed to get data from the paths".into(),
        }
    } else {
        comment = format!(
            "Failed to get the year({:?}) or day({:?}) falling back to default",
            input.year, input.day
        );
    }
    quote! {
        #[doc = #comment]
        const #ident: &str = "";
    }
    .into()
}

/// Same as include_data!, but it only gets data from the example directory.
/// This is useful for testing the examples, even when Data from the AOC is available
#[proc_macro]
pub fn include_example(data: TokenStream) -> TokenStream {
    let input = match parse::<IncludeData>(data) {
        Ok(include) => include,
        Err(error) => return error.into_compile_error().into(),
    };
    let ident = format_ident!("{}", input.var);
    let mut comment: String;
    if input.day.is_some() && input.year.is_some() {
        comment = "Data from the Example dir".into();
        match canonicalize(EXAMPLE_DIR, input.clone()) {
            Some(p) => {
                return quote! {
                    #[doc = #comment]
                    const #ident: &str = include_str!(#p).trim_ascii();
                }
                .into()
            }
            None => comment = "failed to get data from the path".into(),
        }
    } else {
        comment = format!(
            "Failed to get the year({:?}) or day({:?}) falling back to default",
            input.year, input.day
        );
    }
    quote! {
        #[doc = #comment]
        const #ident: &str = "";
    }
    .into()
}