Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
14 kB · 572 lines
Rust
at develop
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573use std::{ borrow::Cow, cmp::Ordering, convert::Infallible, fmt::Display, path::PathBuf, process::ExitCode,};
use pico_args::Arguments;use tree_sitter::{Node, Parser};
const CLI_HELP: &str = r#"USAGE $ rust-organizer [-c] [-w] FILE
ARGUMENTS FILE File name of the Rust source file to reorganize.
FLAGS -c, --check Check whether reorganizing the file would change the file contents. -w, --write Overwrite the file with the reorganized contents."#;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]struct Cli { check: bool, overwrite: bool, path: PathBuf,}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]enum Item<'a> { InnerDoc(Cow<'a, str>), Macro { name: &'a str, content: Cow<'a, str>, }, ModDecl { name: &'a str, content: Cow<'a, str>, }, Use(Cow<'a, str>), Const { name: &'a str, content: Cow<'a, str>, }, Type { name: &'a str, content: Cow<'a, str>, }, Func { name: &'a str, content: Cow<'a, str>, }, Trait { name: &'a str, content: SortableContent<'a>, }, Impl { name: TypeIdent<'a>, trt: Option<&'a str>, content: SortableContent<'a>, }, MacroInvocation(Cow<'a, str>), Mod { name: &'a str, content: SortableContent<'a>, },}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]struct Module<'a> { items: Vec<(bool, Item<'a>)>,}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]struct SortableContent<'a> { before: Cow<'a, str>, inner: Module<'a>, after: Cow<'a, str>,}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]struct TypeIdent<'a> { name: &'a str, generics: Option<&'a str>, reference_type: Option<&'a str>,}
fn main() -> ExitCode { // Parse commandline arguments let mut args = Arguments::from_env(); if args.contains(["-h", "--help"]) { print!("{}", CLI_HELP); return ExitCode::SUCCESS; } let cli: Cli = match args.try_into() { Ok(cli) => cli, Err(e) => { eprintln!("Error: {}", e); return ExitCode::FAILURE; } }; // Run the main program match cli.run() { Ok(code) => code, Err(e) => { eprintln!("Error: {}", e); ExitCode::FAILURE } }}
impl Cli { fn run(&self) -> Result<ExitCode, String> { let mut parser = Parser::new(); parser .set_language(&tree_sitter_rust::LANGUAGE.into()) .expect("Error loading Rust grammar");
let text = std::fs::read_to_string(&self.path) .map_err(|e| format!("unable to read file: {}", e))?;
let Some(tree) = parser.parse(&text, None) else { return Err("unable to parse file".to_owned()); };
let root = tree.root_node(); assert_eq!(root.kind(), "source_file"); let mut root = Module::from_node(&text, root);
let is_sorted = root.is_sorted(self.check); if self.check { return if is_sorted { Ok(ExitCode::SUCCESS) } else { Ok(ExitCode::FAILURE) }; } if self.overwrite && is_sorted { return Ok(ExitCode::SUCCESS); }
root.sort();
if self.overwrite { std::fs::write(&self.path, root.to_string()) .map_err(|e| format!("unable to write file: {}", e))?; } else { println!("{}", root); }
Ok(ExitCode::SUCCESS) }}
impl TryFrom<Arguments> for Cli { type Error = String;
fn try_from(mut args: Arguments) -> Result<Self, Self::Error> { let cli = Cli { check: args.contains(["-c", "--check"]), overwrite: args.contains(["-w", "--write"]), path: args .free_from_os_str::<_, Infallible>(|s| Ok(PathBuf::from(s))) .unwrap(), };
let remaining = args.finish(); match remaining.len() { 0 => Ok(()), 1 => Err(format!( "unexpected argument: '{}'", remaining[0].to_string_lossy() )), _ => Err(format!( "unexpected arguments: {}", remaining .into_iter() .map(|s| format!("'{}'", s.to_string_lossy())) .collect::<Vec<_>>() .join(", ") )), }?; Ok(cli) }}
impl<'a> Item<'a> { fn append_content(&mut self, text: &str) { match self { Item::Macro { content, .. } | Item::ModDecl { content, .. } | Item::Const { content, .. } | Item::Type { content, .. } | Item::Func { content, .. } | Item::InnerDoc(content) | Item::Use(content) | Item::MacroInvocation(content) => { *content = Cow::Owned(format!("{}{}", content, text)); } Item::Impl { .. } | Item::Mod { .. } | Item::Trait { .. } => { // Cannot add content to these items } } }
fn item_order(&self) -> u8 { match self { Item::InnerDoc(_) => 0, Item::Macro { .. } => 1, Item::ModDecl { .. } => 2, Item::Use(_) => 3, Item::Const { .. } => 4, Item::Type { .. } => 5, Item::Trait { .. } => 5, Item::Func { .. } => 6, Item::Impl { .. } => 7, Item::MacroInvocation(_) => 8, Item::Mod { .. } => 9, } }
fn maybe_item(text: &'a str, node: Node<'a>, start: Option<usize>) -> Option<Self> { let get_field_str = |field_name| { node.child_by_field_name(field_name) .map(|n| n.utf8_text(text.as_bytes()).unwrap()) };
let start = start.unwrap_or(node.start_byte()); let end = if node.utf8_text(text.as_bytes()).unwrap().ends_with('\n') { node.end_byte() - 1 } else { node.end_byte() }; let content: Cow<'a, str> = Cow::Borrowed(&text[start..end]); match node.kind() { "attribute_item" => { // Ignore and add to the next item None } "block_comment" | "line_comment" => { let comment = node.utf8_text(text.as_bytes()).unwrap(); if comment.starts_with("//!") || comment.starts_with("/*!") { // Doc comment for the file (ensure that it's at the top of the file). Some(Self::InnerDoc(content)) } else { None // Move comment with the next item } } "const_item" | "static_item" => { let name = get_field_str("name").unwrap(); Some(Self::Const { name, content }) } "associated_type" | "enum_item" | "struct_item" | "type_item" => { let name = get_field_str("name").unwrap(); Some(Self::Type { name, content }) } "function_item" | "function_signature_item" => { let name = get_field_str("name").unwrap(); Some(Self::Func { name, content }) } "trait_item" => { let name = get_field_str("name").unwrap(); let content = SortableContent::within_node(text, node, Some(start), "body"); Some(Self::Trait { name, content }) } "impl_item" => { let name = TypeIdent::from_node(text, node.child_by_field_name("type").unwrap()); let trt = get_field_str("trait"); let content = SortableContent::within_node(text, node, Some(start), "body"); Some(Self::Impl { name, trt, content }) } "inner_attribute_item" => { // Should be at the top (treat like inner doc, to keep it in the chosen // order compared to the module documentation). Some(Self::InnerDoc(content)) } "macro_definition" => { let name = get_field_str("name").unwrap(); Some(Self::Macro { name, content }) } "macro_invocation" => Some(Self::MacroInvocation(content)), "mod_item" => { let name = get_field_str("name").unwrap(); if node.child_by_field_name("body").is_some() { let content = SortableContent::within_node(text, node, Some(start), "body"); Some(Self::Mod { name, content }) } else { Some(Self::ModDecl { name, content }) } } "use_declaration" => Some(Self::Use(content)), _ => panic!( "unexpected node kind: {}\ncontent: {}", node.kind(), content ), } }}
impl Display for Item<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Item::InnerDoc(content) | Item::Macro { content, .. } | Item::MacroInvocation(content) | Item::ModDecl { content, .. } | Item::Use(content) | Item::Const { content, .. } | Item::Type { content, .. } | Item::Func { content, .. } => write!(f, "{content}"), Item::Mod { content, .. } | Item::Impl { content, .. } | Item::Trait { content, .. } => { write!(f, "{content}") } } }}
impl Ord for Item<'_> { fn cmp(&self, other: &Self) -> Ordering { use Item::*;
let self_order = self.item_order(); let other_order = other.item_order(); if self_order != other_order { return self_order.cmp(&other_order); } match (self, other) { (InnerDoc(_), InnerDoc(_)) => Ordering::Equal, (Const { name: a, .. }, Const { name: b, .. }) | (Macro { name: a, .. }, Macro { name: b, .. }) | (Mod { name: a, .. }, Mod { name: b, .. }) | (ModDecl { name: a, .. }, ModDecl { name: b, .. }) | ( Type { name: a, .. } | Trait { name: a, .. }, Type { name: b, .. } | Trait { name: b, .. }, ) | (Func { name: a, .. }, Func { name: b, .. }) => a.cmp(b), (Use(_), Use(_)) | (MacroInvocation(_), MacroInvocation(_)) => Ordering::Equal, ( Impl { name: a, trt: t_a, .. }, Impl { name: b, trt: t_b, .. }, ) => { let name_order = a.name.cmp(b.name); if name_order == Ordering::Equal { let trt_order = t_a.unwrap_or("").cmp(t_b.unwrap_or("")); if trt_order == Ordering::Equal { let a_parts = (a.generics.unwrap_or(""), a.reference_type.unwrap_or("")); let b_parts = (b.generics.unwrap_or(""), b.reference_type.unwrap_or("")); a_parts.cmp(&b_parts) } else { trt_order } } else { name_order } } _ => { // eprintln!("{} -- {}", self, other); unreachable!(); } } }}
impl PartialOrd for Item<'_> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> { Some(self.cmp(other)) }}
impl<'a> Module<'a> { pub fn from_node(text: &'a str, root: Node<'a>) -> Self { assert!(matches!(root.kind(), "source_file" | "declaration_list")); let mut cursor = root.walk(); cursor.goto_first_child();
let mut items: Vec<(bool, Item)> = Vec::new(); let mut start = None; let mut last = None; if cursor.node().kind() == "{" { last = Some(cursor.node().end_byte()); cursor.goto_next_sibling(); } loop { if cursor.node().kind() == "}" { assert!(!cursor.goto_next_sibling()); break; } let node = cursor.node(); // eprintln!("{} : {}\n\n", node.kind(), node.to_sexp()); let inbetween = &text[last.unwrap_or(root.start_byte())..start.unwrap_or(node.start_byte())]; if node.kind() == "empty_statement" { if let Some((_, it)) = items.last_mut() { it.append_content(";"); } debug_assert!( inbetween.trim().is_empty(), "unexpected skipped content: {:?}", inbetween ); start = None; last = Some(node.end_byte()); } else if let Some(item) = Item::maybe_item(text, node, start) { debug_assert!( inbetween.trim().is_empty(), "unexpected skipped content: {:?}", inbetween ); let newline_before = inbetween.contains("\n\n"); items.push((items.is_empty() || newline_before, item)); start = None; last = Some(node.end_byte()); } else if start.is_none() { start = Some(node.start_byte()); } if !cursor.goto_next_sibling() { break; } }
Self { items } }
pub fn is_sorted(&self, print_diff: bool) -> bool { for it in &self.items { match &it.1 { Item::Mod { content, .. } | Item::Impl { content, .. } | Item::Trait { content, .. } => { if !content.is_sorted(print_diff) { return false; } } _ => {} } } for window in self.items.windows(2) { if window[0].1 > window[1].1 { if print_diff { eprintln!( "Expected \n\"\"\"\n{}\n\"\"\"\n before \n\"\"\"\n{}\n\"\"\"", window[1].1, window[0].1 ); } return false; } } true }
pub fn sort(&mut self) { for it in self.items.iter_mut() { match &mut it.1 { Item::Mod { content, .. } | Item::Impl { content, .. } | Item::Trait { content, .. } => content.sort(), _ => {} } } self.items.sort_unstable_by(|a, b| a.1.cmp(&b.1)); }}
impl Display for Module<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let mut last = None; for (newline, item) in &self.items { if *newline || (last.is_some() && last != Some(item.item_order())) { writeln!(f)?; } writeln!(f, "{}", item)?; last = Some(item.item_order()); } Ok(()) }}
impl<'a> SortableContent<'a> { fn is_sorted(&self, print_diff: bool) -> bool { self.inner.is_sorted(print_diff) }
fn sort(&mut self) { self.inner.sort(); }
fn within_node( text: &'a str, node: Node<'a>, start: Option<usize>, child: &'static str, ) -> Self { let start = start.unwrap_or(node.start_byte()); let body = node.child_by_field_name(child).unwrap();
let mut cursor = body.walk(); cursor.goto_first_child(); assert_eq!(cursor.node().kind(), "{"); let before = Cow::Borrowed(&text[start..cursor.node().end_byte()]);
cursor.goto_parent(); cursor.goto_last_child(); assert_eq!(cursor.node().kind(), "}"); let after = Cow::Borrowed(&text[cursor.node().start_byte()..node.end_byte()]);
let inner = Module::from_node(text, body); Self { before, inner, after, } }}
impl Display for SortableContent<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}{}{}", self.before, self.inner, self.after) }}
impl<'a> TypeIdent<'a> { fn from_node(text: &'a str, node: Node<'a>) -> Self { let get_field_str = |field_name| { node.child_by_field_name(field_name) .map(|n| n.utf8_text(text.as_bytes()).unwrap()) };
match node.kind() { "array_type" => { let inner = node.child_by_field_name("element").unwrap(); let mut ty = TypeIdent::from_node(text, inner); let reference_str = &text[node.start_byte()..inner.start_byte()]; ty.reference_type = Some(reference_str); ty } "generic_type" => { let name = get_field_str("type").unwrap(); let generics = get_field_str("type_arguments"); debug_assert!(generics.is_some()); Self { name, generics, reference_type: None, } } "reference_type" => { let inner = node.child_by_field_name("type").unwrap(); let mut ty = TypeIdent::from_node(text, inner); let reference_str = &text[node.start_byte()..inner.start_byte()]; ty.reference_type = Some(reference_str); ty } "type_identifier" | "scoped_type_identifier" | "primitive_type" | "bounded_type" => { Self { name: node.utf8_text(text.as_bytes()).unwrap(), generics: None, reference_type: None, } } _ => panic!( "invalid type identifier node: {}, `{}'", node.kind(), node.utf8_text(text.as_bytes()).unwrap() ), } }}