use std::collections::hash_map::Entry;
use crate::{
hash::{HashMap, HashSet},
options::Options,
parser::{FileScope, ModuleContext, QualifiedPath},
};
#[derive(Debug, Clone)]
pub struct RewriteResult {
pub content: String,
pub replacements: Vec<(String, String)>,
}
#[derive(Debug)]
enum Action {
Insert {
pos: usize,
text: String,
},
Replace {
start: usize,
end: usize,
original: String,
replacement: String,
},
}
impl Action {
#[inline]
fn pos(&self) -> usize {
match self {
Action::Insert { pos, .. } => *pos,
Action::Replace { start, .. } => *start,
}
}
}
pub fn compute_line_offsets(source: &str) -> Vec<usize> {
let mut offsets = Vec::with_capacity(source.len() / 40 + 1);
offsets.push(0);
for idx in memchr::memchr_iter(b'\n', source.as_bytes()) {
offsets.push(idx + 1);
}
offsets
}
pub fn rewrite_source(
source: &str,
paths: &[QualifiedPath],
insert_pos: usize,
scope: &FileScope,
options: &Options,
) -> Option<RewriteResult> {
let line_offsets = compute_line_offsets(source);
rewrite_source_with_offsets(source, &line_offsets, paths, insert_pos, scope, options)
}
pub fn rewrite_source_with_offsets(
source: &str,
line_offsets: &[usize],
paths: &[QualifiedPath],
insert_pos: usize,
scope: &FileScope,
options: &Options,
) -> Option<RewriteResult> {
let root_mod = ModuleContext {
insert_pos,
indent: String::new(),
has_existing_use: !scope.existing_use_paths.is_empty(),
scope: scope.clone(),
paths: paths.to_vec(),
};
rewrite_modules_with_offsets(source, line_offsets, &[root_mod], options)
}
pub fn rewrite_modules_with_offsets(
source: &str,
line_offsets: &[usize],
modules: &[ModuleContext],
_options: &Options,
) -> Option<RewriteResult> {
if modules.is_empty() {
return None;
}
let mut actions = Vec::new();
for module in modules {
if module.paths.is_empty() {
continue;
}
let paths = &module.paths;
let scope = &module.scope;
let mut needed_imports = HashSet::default();
let mut new_in_scope = HashSet::default();
let is_in_scope = |ident: &str, new_in_scope: &HashSet<String>| {
scope.in_scope_idents.contains(ident) || new_in_scope.contains(ident)
};
let conflicting_tails: HashSet<&str> = if paths.len() > 1 {
let mut tail_to_import: HashMap<&str, &str> = HashMap::default();
let mut conflicts = HashSet::default();
for qp in paths {
let tail = qp.last_ident();
let import = qp.import.as_str();
match tail_to_import.entry(tail) {
Entry::Vacant(e) => {
e.insert(import);
}
Entry::Occupied(e) => {
if *e.get() != import {
conflicts.insert(tail);
}
}
}
}
conflicts
} else {
HashSet::default()
};
for qp in paths {
let span_start = qp.span.start();
if span_start.line == 0 || span_start.line > line_offsets.len() {
continue;
}
let line_start = line_offsets[span_start.line - 1];
let line_end = if span_start.line < line_offsets.len() {
line_offsets[span_start.line]
} else {
source.len()
};
let line_text = &source[line_start..line_end];
let target = &qp.original;
let trimmed_target = target.trim_start_matches(':');
let start_search = line_text
.char_indices()
.nth(span_start.column)
.map(|(idx, _)| idx)
.unwrap_or(span_start.column.min(line_text.len()));
let match_info = if line_text[start_search..].starts_with(target) {
Some((start_search, target.len()))
} else if line_text[start_search..].starts_with(trimmed_target) {
Some((start_search, trimmed_target.len()))
} else if let Some(pos) = line_text[start_search..].find(target) {
Some((start_search + pos, target.len()))
} else if let Some(pos) = line_text[start_search..].find(trimmed_target) {
Some((start_search + pos, trimmed_target.len()))
} else if let Some(pos) = line_text.find(target) {
Some((pos, target.len()))
} else {
line_text
.find(trimmed_target)
.map(|pos| (pos, trimmed_target.len()))
};
let Some((col_start, matched_len)) = match_info else {
continue;
};
let start_byte = line_start + col_start;
let end_byte = start_byte + matched_len;
if end_byte <= source.len() {
let mut replacement = qp.replacement.clone();
let mut import = qp.import.clone();
let total = qp.segments.len();
if conflicting_tails.contains(qp.last_ident()) && total >= 2 {
let parent_mod = &qp.segments[total - 2];
if is_in_scope(parent_mod.as_str(), &new_in_scope) {
continue;
}
import = qp.segments[..total - 1].join("::");
replacement = qp.segments[total - 2..].join("::");
}
let import_tail = import.rsplit("::").next().unwrap_or_default();
if scope.existing_use_paths.contains(&import) || needed_imports.contains(&import) {
} else if is_in_scope(import_tail, &new_in_scope) {
if total >= 2 {
let parent_mod = &qp.segments[total - 2];
let candidate_import = qp.segments[..total - 1].join("::");
let candidate_tail = parent_mod.as_str();
if scope.existing_use_paths.contains(&candidate_import)
|| needed_imports.contains(&candidate_import)
{
replacement = qp.segments[total - 2..].join("::");
} else if !is_in_scope(candidate_tail, &new_in_scope) {
import = candidate_import;
replacement = qp.segments[total - 2..].join("::");
needed_imports.insert(import.clone());
new_in_scope.insert(candidate_tail.to_string());
} else {
continue;
}
} else {
continue;
}
} else {
needed_imports.insert(import.clone());
new_in_scope.insert(import_tail.to_string());
}
actions.push(Action::Replace {
start: start_byte,
end: end_byte,
original: qp.original.clone(),
replacement,
});
}
}
let mut to_insert: Vec<String> = needed_imports
.into_iter()
.filter(|imp| !scope.existing_use_paths.contains(imp))
.collect();
if !to_insert.is_empty() {
to_insert.sort_unstable();
let indent_len = module.indent.len();
let total_len: usize = to_insert.iter().map(|s| indent_len + s.len() + 6).sum::<usize>() + 2;
let mut use_block = String::with_capacity(total_len);
for imp in &to_insert {
use_block.push_str(&module.indent);
use_block.push_str("use ");
use_block.push_str(imp);
use_block.push_str(";\n");
}
if !module.has_existing_use
&& module.insert_pos < source.len()
&& !source[module.insert_pos..].starts_with('\n')
{
use_block.push('\n');
}
actions.push(Action::Insert {
pos: module.insert_pos,
text: use_block,
});
}
}
if actions.is_empty() {
return None;
}
actions.sort_unstable_by(|a, b| {
a.pos().cmp(&b.pos()).then_with(|| match (a, b) {
(Action::Insert { .. }, Action::Replace { .. }) => std::cmp::Ordering::Less,
(Action::Replace { .. }, Action::Insert { .. }) => std::cmp::Ordering::Greater,
(Action::Replace { end: end_a, .. }, Action::Replace { end: end_b, .. }) => {
end_b.cmp(end_a)
}
_ => std::cmp::Ordering::Equal,
})
});
let mut last_end = 0;
actions.retain(|act| match act {
Action::Insert { .. } => true,
Action::Replace { start, end, .. } => {
if *start >= last_end {
last_end = *end;
true
} else {
false
}
}
});
let mut current = String::with_capacity(source.len() + 128);
let mut cur_idx = 0;
let mut replacements = Vec::with_capacity(actions.len());
for act in actions {
match act {
Action::Insert { pos, text } => {
if pos > cur_idx {
current.push_str(&source[cur_idx..pos]);
cur_idx = pos;
}
current.push_str(&text);
}
Action::Replace {
start,
end,
original,
replacement,
} => {
if start > cur_idx {
current.push_str(&source[cur_idx..start]);
}
current.push_str(&replacement);
cur_idx = end;
replacements.push((original, replacement));
}
}
}
current.push_str(&source[cur_idx..]);
if current == source {
return None;
}
Some(RewriteResult {
content: current,
replacements,
})
}