1use core::str;
2use std::{
3    fs, io,
4    ops::Range,
5    path::{Path, PathBuf},
6    str::FromStr,
7};
8
9use annotate_snippets::{Level, Renderer, Snippet};
10use anyhow::anyhow;
11use proc_macro2::{Span, TokenStream};
12use syn::{
13    parse2, spanned::Spanned, visit::Visit, Attribute, Expr, ExprLit, File, ItemMod, Lit, Meta,
14};
15
16pub struct Source {
17    path: PathBuf,
18    text: String,
19}
20
21pub struct ExpandError {
22    e: anyhow::Error,
23    span: Option<Range<usize>>,
24    source: Option<Source>,
25}
26impl ExpandError {
27    fn new(span: Option<Span>, e: impl Into<anyhow::Error>) -> Self {
28        let e = e.into();
29        let span = span.map(|s| s.byte_range());
30        let source = None;
31        Self { e, span, source }
32    }
33    pub fn show(&self) {
34        let title = self.e.to_string();
35        let path;
36        let mut m = Level::Error.title(&title);
37        if let (Some(source), Some(span)) = (&self.source, self.span.clone()) {
38            path = source.path.to_string_lossy();
39            m = m.snippet(
40                Snippet::source(&source.text)
41                    .fold(true)
42                    .origin(&path)
43                    .annotation(Level::Error.span(span)),
44            );
45        }
46        let renderer = Renderer::styled();
47        eprintln!("{}", renderer.render(m));
48    }
49}
50
51impl<E: Into<anyhow::Error>> From<E> for ExpandError {
52    fn from(e: E) -> Self {
53        let e = e.into();
54        let span = e.downcast_ref::<syn::Error>().map(|e| e.span());
55        Self::new(span, e)
56    }
57}
58type Result<T> = std::result::Result<T, ExpandError>;
59
60fn with_path<T>(r: io::Result<T>, path: &Path) -> Result<T> {
61    r.map_err(|e| anyhow!("Could not read file : `{}` ({e})", path.display()).into())
62}
63fn with_source<T>(r: Result<T>, path: &Path, text: &str) -> Result<T> {
64    r.map_err(|mut e| {
65        if e.source.is_none() {
66            let path = path.to_path_buf();
67            let text = text.to_string();
68            e.source = Some(Source { path, text });
69        }
70        e
71    })
72}
73
74pub fn expand_from_path(root: &Path, path: &Path, is_root: bool) -> Result<String> {
75    if path.canonicalize()?.strip_prefix(root).is_err() {
76        return Err(ExpandError::new(
77            None,
78            anyhow!("path is out of root directory : `{}`", path.display()),
79        ));
80    }
81
82    let s = with_path(fs::read_to_string(path), path)?;
83    with_source(expand_from_text(root, path, is_root, &s), path, &s)
84}
85fn expand_from_text(root: &Path, path: &Path, is_root: bool, s: &str) -> Result<String> {
86    let tokens = parse_token_stream(s)?;
87    let file: File = parse2(tokens)?;
88    let mut b = CodeBuilder::new();
89    b.visit_file(&file);
90    let mut text = String::new();
91    for part in b.finish(s.len()) {
92        match part {
93            Part::Text(r) => text.push_str(&s[r]),
94            Part::Mod(m) => {
95                text.push_str(" {\n");
96                text.push_str(&expand_from_path(
97                    root,
98                    &path_from_mod(path, is_root, &m)?,
99                    false,
100                )?);
101                text.push_str("}\n");
102            }
103        }
104    }
105    Ok(text)
106}
107
108fn parse_token_stream(s: &str) -> syn::Result<TokenStream> {
109    match TokenStream::from_str(s) {
110        Ok(tokens) => Ok(tokens),
111        Err(e) => Err(syn::Error::new(e.span(), e)),
112    }
113}
114
115fn path_from_mod(path: &Path, is_root: bool, m: &ItemMod) -> Result<PathBuf> {
116    match path_from_attrs(&m.attrs) {
117        Some(p) => Ok(path.parent().unwrap().join(p)),
118        None => {
119            let name = m.ident.to_string();
120            let file_name = path.file_name().unwrap();
121            let base = if is_root || file_name == "mod.rs" {
122                path.parent().unwrap().to_path_buf()
123            } else {
124                path.with_extension("")
125            };
126            let p0 = base.join(format!("{name}.rs"));
127            let p1 = base.join(format!("{name}/mod.rs"));
128            for p in &[&p0, &p1] {
129                if p.is_file() {
130                    return Ok(p.to_path_buf());
131                }
132            }
133            Err(ExpandError::new(
134                Some(m.span()),
135                anyhow!("Could not find source file : `{}`", p0.display()),
136            ))
137        }
138    }
139}
140
141fn path_from_attrs(attr: &[Attribute]) -> Option<PathBuf> {
142    for attr in attr {
143        if let Some(p) = path_from_attr(attr) {
144            return Some(p);
145        }
146    }
147    None
148}
149
150fn path_from_attr(attr: &Attribute) -> Option<PathBuf> {
151    match &attr.meta {
152        Meta::NameValue(nv) => {
153            if nv.path.is_ident("path") {
154                if let Expr::Lit(ExprLit {
155                    lit: Lit::Str(s), ..
156                }) = &nv.value
157                {
158                    return Some(PathBuf::from(s.value()));
159                }
160            }
161            None
162        }
163        _ => None,
164    }
165}
166
167enum Part {
168    Text(Range<usize>),
169    Mod(ItemMod),
170}
171
172struct CodeBuilder {
173    offset: usize,
174    parts: Vec<Part>,
175}
176impl CodeBuilder {
177    fn new() -> Self {
178        Self {
179            offset: 0,
180            parts: Vec::new(),
181        }
182    }
183    fn finish(self, source_len: usize) -> Vec<Part> {
184        let mut parts = self.parts;
185        parts.push(Part::Text(self.offset..source_len));
186        parts
187    }
188}
189impl<'ast> Visit<'ast> for CodeBuilder {
190    fn visit_item_mod(&mut self, i: &'ast ItemMod) {
191        if i.content.is_some() {
192            return;
193        }
194        let end = i.ident.span().byte_range().end;
195        self.parts.push(Part::Text(self.offset..end));
196        self.parts.push(Part::Mod(i.clone()));
197        self.offset = i.span().byte_range().end;
198    }
199}