Skip to main content

krabka_protocol_codegen/
hoist.rs

1//! Hoist absolute paths in generated code into module-scope `use` items.
2//!
3//! The emitters name an item from another module by its full path, such as
4//! `crate::primitives::array::put_array_len` or
5//! `crate::owned::fetch_request::PartitionData`. That keeps each emitter local:
6//! it never has to know which names the rest of the file imports. The
7//! workspace denies `clippy::absolute_paths` beyond two segments, though, so
8//! the committed code must name such an item through an import.
9//!
10//! [`hoist_absolute_paths`] closes that gap on the parsed file. It finds every
11//! path that starts at a crate root and names an item through three or more
12//! segments, as the lint counts them, picks the
13//! deepest segment that it can import without a name clash, adds one `use`
14//! item per imported path, and shortens each path to start at the imported
15//! name. A name that two different paths would import, or that the file
16//! already binds, falls back to its parent module, so
17//! `crate::owned::fetch_request::PartitionData` in a file that defines its own
18//! `PartitionData` becomes `fetch_request::PartitionData`.
19//!
20//! A generated file is `include!`d into its wrapper module, so its module-scope
21//! `use` items are that module's imports.
22
23use std::{
24    collections::{BTreeMap, BTreeSet},
25    mem,
26};
27
28use proc_macro2::Span;
29use quote::ToTokens;
30use syn::{
31    Expr, ExprPath, File, Ident, Item, ItemUse, Macro, PatIdent, Path, Token, TypeParam, TypePath,
32    UseName, UsePath, UseTree, Visibility,
33    punctuated::Punctuated,
34    visit::{self, Visit},
35    visit_mut::{self, VisitMut},
36};
37
38/// The longest crate-rooted path that `clippy::absolute_paths` accepts, as
39/// `clippy.toml` sets `absolute-paths-max-segments`.
40const MAX_SEGMENTS: usize = 2;
41
42/// The crate roots whose paths count as absolute even when nothing in the file
43/// names them: the standard library facades.
44const STD_ROOTS: [&str; 3] = ["std", "core", "alloc"];
45
46/// The standard-prelude names a generated file uses through a full path, with
47/// the path below the `std` root that each one names.
48///
49/// Such a name needs no `use` item, and no other path may import it.
50const PRELUDE: [(&str, &[&str]); 14] = [
51    ("Box", &["boxed", "Box"]),
52    ("Clone", &["clone", "Clone"]),
53    ("Default", &["default", "Default"]),
54    ("Err", &["result", "Result", "Err"]),
55    ("From", &["convert", "From"]),
56    ("Into", &["convert", "Into"]),
57    ("None", &["option", "Option", "None"]),
58    ("Ok", &["result", "Result", "Ok"]),
59    ("Option", &["option", "Option"]),
60    ("Result", &["result", "Result"]),
61    ("Some", &["option", "Option", "Some"]),
62    ("String", &["string", "String"]),
63    ("ToString", &["string", "ToString"]),
64    ("Vec", &["vec", "Vec"]),
65];
66
67/// A path, keyed by its leading `::` and its segment names, without generic
68/// arguments.
69type PathKey = (bool, Vec<String>);
70
71/// Shorten every crate-rooted path in `file` that names an item through three
72/// or more segments to start at an imported name, and add the `use` items that
73/// import those names.
74///
75/// The function also shortens the paths in a macro call whose body is a list
76/// of expressions, such as `vec![..]`. It leaves `use` items, visibility
77/// paths, qualified paths such as `<T as Trait>::f`, and every other macro
78/// body alone. A path that it cannot shorten without a clash at every level
79/// stays as it is.
80pub fn hoist_absolute_paths(file: &mut File) {
81    let scope = Scope::of(file);
82    let mut collector = Collector {
83        scope: &scope,
84        paths: BTreeSet::new(),
85    };
86    collector.visit_file(file);
87    let levels = choose_levels(&scope, &collector.paths);
88    if levels.is_empty() {
89        return;
90    }
91
92    let imports: BTreeSet<PathKey> = levels
93        .iter()
94        .map(|((colon, segs), &level)| (*colon, segs.iter().take(level + 1).cloned().collect()))
95        .filter(|import| !scope.already_imports(import) && !is_prelude(import))
96        .collect();
97
98    Rewriter {
99        scope: &scope,
100        levels: &levels,
101        rewritten: 0,
102    }
103    .visit_file_mut(file);
104
105    let at = file
106        .items
107        .iter()
108        .rposition(|item| matches!(item, Item::Use(_)))
109        .map_or(0, |last| last + 1);
110    let new_items = imports.iter().filter_map(use_item).map(Item::Use);
111    file.items.splice(at..at, new_items);
112}
113
114/// The names a file binds, split by where they bind them.
115#[derive(Debug, Default)]
116struct Scope {
117    /// Module-scope `use` items: the bound name and the path it imports.
118    imported: BTreeMap<String, PathKey>,
119    /// Names in the type namespace that a path segment could resolve to: the
120    /// module-scope items, every nested item and `use`, and generic parameters.
121    types: BTreeSet<String>,
122    /// Names in the value namespace bound by a pattern, such as a local
123    /// variable or a parameter.
124    values: BTreeSet<String>,
125}
126
127impl Scope {
128    fn of(file: &File) -> Self {
129        let mut scope = Self::default();
130        for item in &file.items {
131            if let Item::Use(item_use) = item {
132                use_bindings(
133                    &item_use.tree,
134                    &mut Vec::new(),
135                    item_use.leading_colon.is_some(),
136                    &mut |name, path| {
137                        scope.imported.insert(name, path);
138                    },
139                );
140            }
141        }
142        let mut binder = Binder {
143            scope: &mut scope,
144            depth: 0,
145        };
146        binder.visit_file(file);
147        scope
148    }
149
150    /// True when `name` is bound in the file other than by a module-scope
151    /// `use` item.
152    fn binds_type(&self, name: &str) -> bool {
153        self.types.contains(name)
154    }
155
156    /// True when `path` starts at a crate root rather than at a name the file
157    /// binds.
158    fn is_crate_rooted(&self, path: &Path) -> bool {
159        let Some(first) = path.segments.first() else {
160            return false;
161        };
162        let name = first.ident.to_string();
163        if path.leading_colon.is_some() || name == "crate" || STD_ROOTS.contains(&name.as_str()) {
164            return true;
165        }
166        let starts_lowercase = name.starts_with(|c: char| c.is_ascii_lowercase());
167        starts_lowercase
168            && name != "self"
169            && name != "super"
170            && !self.imported.contains_key(&name)
171            && !self.binds_type(&name)
172    }
173
174    /// True when a module-scope `use` item already imports exactly `import`.
175    fn already_imports(&self, import: &PathKey) -> bool {
176        import
177            .1
178            .last()
179            .and_then(|name| self.imported.get(name))
180            .is_some_and(|bound| bound == import)
181    }
182}
183
184/// Walks a file and records the names it binds outside its module-scope `use`
185/// items.
186struct Binder<'s> {
187    scope: &'s mut Scope,
188    /// How many items deep the walk is. A `use` at depth zero is a
189    /// module-scope import, which [`Scope::of`] records separately.
190    depth: usize,
191}
192
193impl Binder<'_> {
194    fn bind_item(&mut self, item: &Item) {
195        let ident = match item {
196            Item::Const(i) => Some(&i.ident),
197            Item::Enum(i) => Some(&i.ident),
198            Item::ExternCrate(i) => Some(i.rename.as_ref().map_or(&i.ident, |(_, rename)| rename)),
199            Item::Fn(i) => Some(&i.sig.ident),
200            Item::Macro(i) => i.ident.as_ref(),
201            Item::Mod(i) => Some(&i.ident),
202            Item::Static(i) => Some(&i.ident),
203            Item::Struct(i) => Some(&i.ident),
204            Item::Trait(i) => Some(&i.ident),
205            Item::TraitAlias(i) => Some(&i.ident),
206            Item::Type(i) => Some(&i.ident),
207            Item::Union(i) => Some(&i.ident),
208            Item::Use(i) if self.depth > 0 => {
209                let types = &mut self.scope.types;
210                use_bindings(
211                    &i.tree,
212                    &mut Vec::new(),
213                    i.leading_colon.is_some(),
214                    &mut |name, _| {
215                        types.insert(name);
216                    },
217                );
218                None
219            }
220            _ => None,
221        };
222        if let Some(ident) = ident {
223            self.scope.types.insert(ident.to_string());
224        }
225    }
226}
227
228impl<'ast> Visit<'ast> for Binder<'_> {
229    fn visit_item(&mut self, item: &'ast Item) {
230        self.bind_item(item);
231        self.depth += 1;
232        visit::visit_item(self, item);
233        self.depth -= 1;
234    }
235
236    fn visit_pat_ident(&mut self, pat: &'ast PatIdent) {
237        self.scope.values.insert(pat.ident.to_string());
238        visit::visit_pat_ident(self, pat);
239    }
240
241    fn visit_type_param(&mut self, param: &'ast TypeParam) {
242        self.scope.types.insert(param.ident.to_string());
243        visit::visit_type_param(self, param);
244    }
245}
246
247/// Call `bind` with each name a `use` tree binds and the path it imports. A
248/// glob binds no single name, so the function skips it.
249fn use_bindings(
250    tree: &UseTree,
251    prefix: &mut Vec<String>,
252    leading_colon: bool,
253    bind: &mut impl FnMut(String, PathKey),
254) {
255    match tree {
256        UseTree::Path(p) => {
257            prefix.push(p.ident.to_string());
258            use_bindings(&p.tree, prefix, leading_colon, bind);
259            prefix.pop();
260        }
261        UseTree::Name(n) => {
262            let mut path = prefix.clone();
263            path.push(n.ident.to_string());
264            bind(n.ident.to_string(), (leading_colon, path));
265        }
266        UseTree::Rename(r) => {
267            let mut path = prefix.clone();
268            path.push(r.ident.to_string());
269            bind(r.rename.to_string(), (leading_colon, path));
270        }
271        UseTree::Group(g) => {
272            for tree in &g.items {
273                use_bindings(tree, prefix, leading_colon, bind);
274            }
275        }
276        UseTree::Glob(_) => {}
277    }
278}
279
280/// The key of `path`: its leading `::` and its segment names.
281fn key_of(path: &Path) -> PathKey {
282    (
283        path.leading_colon.is_some(),
284        path.segments.iter().map(|s| s.ident.to_string()).collect(),
285    )
286}
287
288/// True when `path` is crate-rooted and names an item through more segments
289/// than `clippy::absolute_paths` accepts.
290///
291/// Like the lint, this function measures only the path to the item. A member
292/// of a type, such as `bytes::Bytes::new`, counts as the path to its type.
293fn is_hoistable(scope: &Scope, path: &Path) -> bool {
294    let (_, segs) = key_of(path);
295    // The path to the item runs through segment `deepest_level`, so it has one
296    // segment more than that index.
297    deepest_level(&segs) >= MAX_SEGMENTS && scope.is_crate_rooted(path)
298}
299
300/// The arguments of a macro call whose body is a comma-separated list of
301/// expressions, such as `vec![..]` or `assert!(..)`. Any other body, such as a
302/// `json!` object, fails to parse, and the walk leaves it alone.
303fn macro_args(mac: &Macro) -> syn::Result<Punctuated<Expr, Token![,]>> {
304    mac.parse_body_with(Punctuated::parse_terminated)
305}
306
307/// Collects the key of every hoistable path in a file.
308struct Collector<'s> {
309    scope: &'s Scope,
310    paths: BTreeSet<PathKey>,
311}
312
313impl<'ast> Visit<'ast> for Collector<'_> {
314    fn visit_item_use(&mut self, _: &'ast ItemUse) {}
315
316    fn visit_visibility(&mut self, _: &'ast Visibility) {}
317
318    fn visit_expr_path(&mut self, expr: &'ast ExprPath) {
319        match &expr.qself {
320            Some(qself) => self.visit_type(&qself.ty),
321            None => visit::visit_expr_path(self, expr),
322        }
323    }
324
325    fn visit_type_path(&mut self, ty: &'ast TypePath) {
326        match &ty.qself {
327            Some(qself) => self.visit_type(&qself.ty),
328            None => visit::visit_type_path(self, ty),
329        }
330    }
331
332    fn visit_macro(&mut self, mac: &'ast Macro) {
333        if let Ok(args) = macro_args(mac) {
334            for arg in &args {
335                self.visit_expr(arg);
336            }
337        }
338    }
339
340    fn visit_path(&mut self, path: &'ast Path) {
341        if is_hoistable(self.scope, path) {
342            self.paths.insert(key_of(path));
343        }
344        visit::visit_path(self, path);
345    }
346}
347
348/// Rewrites every hoisted path to start at its imported name.
349struct Rewriter<'s> {
350    scope: &'s Scope,
351    levels: &'s BTreeMap<PathKey, usize>,
352    /// How many paths the walk has shortened so far.
353    rewritten: usize,
354}
355
356impl VisitMut for Rewriter<'_> {
357    fn visit_item_use_mut(&mut self, _: &mut ItemUse) {}
358
359    fn visit_visibility_mut(&mut self, _: &mut Visibility) {}
360
361    fn visit_expr_path_mut(&mut self, expr: &mut ExprPath) {
362        match &mut expr.qself {
363            Some(qself) => self.visit_type_mut(&mut qself.ty),
364            None => visit_mut::visit_expr_path_mut(self, expr),
365        }
366    }
367
368    fn visit_type_path_mut(&mut self, ty: &mut TypePath) {
369        match &mut ty.qself {
370            Some(qself) => self.visit_type_mut(&mut qself.ty),
371            None => visit_mut::visit_type_path_mut(self, ty),
372        }
373    }
374
375    fn visit_macro_mut(&mut self, mac: &mut Macro) {
376        let Ok(mut args) = macro_args(mac) else {
377            return;
378        };
379        let before = self.rewritten;
380        for arg in &mut args {
381            self.visit_expr_mut(arg);
382        }
383        if self.rewritten != before {
384            mac.tokens = args.into_token_stream();
385        }
386    }
387
388    fn visit_path_mut(&mut self, path: &mut Path) {
389        if is_hoistable(self.scope, path)
390            && let Some(&level) = self.levels.get(&key_of(path))
391        {
392            let segments = mem::take(&mut path.segments);
393            path.segments = segments.into_pairs().skip(level).collect();
394            path.leading_colon = None;
395            self.rewritten += 1;
396        }
397        visit_mut::visit_path_mut(self, path);
398    }
399}
400
401/// The deepest segment index a path may import at: its last segment, or its
402/// first type-like (capitalised) segment, because the members of a type or
403/// trait cannot be imported on their own.
404fn deepest_level(segs: &[String]) -> usize {
405    let last = segs.len().saturating_sub(1);
406    segs.iter()
407        .enumerate()
408        .skip(1)
409        .find(|(_, s)| s.starts_with(|c: char| c.is_ascii_uppercase()))
410        .map_or(last, |(i, _)| i)
411}
412
413/// The prelude path a name stands for, if it is a prelude name.
414fn prelude_path(name: &str) -> Option<&'static [&'static str]> {
415    PRELUDE
416        .iter()
417        .find_map(|(prelude, path)| (*prelude == name).then_some(*path))
418}
419
420/// True when `import` names a prelude item through the standard library, so
421/// the item is already in scope.
422fn is_prelude(import: &PathKey) -> bool {
423    let Some((root, rest)) = import.1.split_first() else {
424        return false;
425    };
426    STD_ROOTS.contains(&root.as_str())
427        && rest
428            .last()
429            .and_then(|name| prelude_path(name))
430            .is_some_and(|path| path.iter().copied().eq(rest.iter().map(String::as_str)))
431}
432
433/// Pick the segment each hoistable path imports at, as an index into its
434/// segments. A path that clashes at every level is left out.
435///
436/// Every path starts at its [`deepest_level`]. A path moves one level up while
437/// the name it would import clashes: the file binds it, a module-scope `use`
438/// binds it to a different path, it is a prelude name for a different path,
439/// or another path would import the same name from a different place. Levels
440/// only fall, so the loop ends.
441fn choose_levels(scope: &Scope, paths: &BTreeSet<PathKey>) -> BTreeMap<PathKey, usize> {
442    let mut levels: BTreeMap<PathKey, usize> = paths
443        .iter()
444        .map(|p| (p.clone(), deepest_level(&p.1)))
445        .collect();
446    let mut changed = true;
447    while changed {
448        changed = false;
449        let claims = claims(&levels);
450        let mut dropped = Vec::new();
451        for (path, level) in &mut levels {
452            let Some((import, name)) = import_at(path, *level) else {
453                dropped.push(path.clone());
454                continue;
455            };
456            if clashes(scope, &claims, &import, name, *level + 1 < path.1.len()) {
457                changed = true;
458                if *level > 1 {
459                    *level -= 1;
460                } else {
461                    dropped.push(path.clone());
462                }
463            }
464        }
465        for path in dropped {
466            levels.remove(&path);
467        }
468    }
469    levels
470}
471
472/// The import a path takes at `level`, and the name it binds.
473fn import_at(path: &PathKey, level: usize) -> Option<(PathKey, &str)> {
474    let name = path.1.get(level)?;
475    let import = (path.0, path.1.iter().take(level + 1).cloned().collect());
476    Some((import, name.as_str()))
477}
478
479/// Every name the current levels import, with each distinct path that imports
480/// it.
481fn claims(levels: &BTreeMap<PathKey, usize>) -> BTreeMap<String, BTreeSet<PathKey>> {
482    let mut claims: BTreeMap<String, BTreeSet<PathKey>> = BTreeMap::new();
483    for (path, &level) in levels {
484        if let Some((import, name)) = import_at(path, level) {
485            claims.entry(name.to_owned()).or_default().insert(import);
486        }
487    }
488    claims
489}
490
491/// True when importing `import` under `name` would clash with another binding.
492///
493/// `prefix_only` says the name only ever opens a longer path. Such a name
494/// resolves in the type namespace, so a local variable of the same name does
495/// not shadow it.
496fn clashes(
497    scope: &Scope,
498    claims: &BTreeMap<String, BTreeSet<PathKey>>,
499    import: &PathKey,
500    name: &str,
501    prefix_only: bool,
502) -> bool {
503    if scope.binds_type(name) || (!prefix_only && scope.values.contains(name)) {
504        return true;
505    }
506    if scope
507        .imported
508        .get(name)
509        .is_some_and(|bound| bound != import)
510    {
511        return true;
512    }
513    if prelude_path(name).is_some() && !is_prelude(import) {
514        return true;
515    }
516    claims.get(name).is_some_and(|paths| paths.len() > 1)
517}
518
519/// A `use` item that imports `import`, or `None` for an empty path.
520fn use_item(import: &PathKey) -> Option<ItemUse> {
521    let (leading_colon, segs) = import;
522    let (name, parents) = segs.split_last()?;
523    let mut tree = UseTree::Name(UseName {
524        ident: Ident::new(name, Span::call_site()),
525    });
526    for parent in parents.iter().rev() {
527        tree = UseTree::Path(UsePath {
528            ident: Ident::new(parent, Span::call_site()),
529            colon2_token: <Token![::]>::default(),
530            tree: Box::new(tree),
531        });
532    }
533    Some(ItemUse {
534        attrs: Vec::new(),
535        vis: Visibility::Inherited,
536        use_token: <Token![use]>::default(),
537        leading_colon: leading_colon.then(<Token![::]>::default),
538        tree,
539        semi_token: <Token![;]>::default(),
540    })
541}