use std::convert::Infallible;
use eure_tree::prelude::*;
use rand::prelude::*;
pub fn unformat(tree: &mut Cst) {
let mut unformatter = Unformatter::new();
let Ok(_) = tree.visit_from_root(&mut unformatter);
unformatter.commands.apply_to(tree).unwrap();
}
pub fn unformat_with_seed(tree: &mut Cst, seed: u64) {
let mut unformatter = Unformatter {
rng: rand::rngs::SmallRng::seed_from_u64(seed),
..Default::default()
};
let Ok(_) = tree.visit_from_root(&mut unformatter);
unformatter.commands.apply_to(tree).unwrap();
}
pub struct Unformatter {
commands: CstCommands,
rng: rand::rngs::SmallRng,
no_white_space_here: bool,
weird_space_probability: f32,
empty_line_probability: f32,
line_removal_probability: f32,
whitespace_removal_probability: f32,
parent: Option<CstNodeId>,
}
impl Unformatter {
pub fn new() -> Self {
Self {
commands: CstCommands::default(),
rng: rand::rngs::SmallRng::from_os_rng(),
no_white_space_here: false,
weird_space_probability: 0.2,
empty_line_probability: 0.2,
line_removal_probability: 0.2,
whitespace_removal_probability: 0.2,
parent: None,
}
}
fn random_whitespace(&mut self) -> String {
let ws_kind = [' ', '\t', '\u{3000}'];
let ws_kind = ws_kind[self.rng.random_range(0..ws_kind.len())];
ws_kind.to_string().repeat(self.rng.random_range(1..=5))
}
}
impl Default for Unformatter {
fn default() -> Self {
Self::new()
}
}
impl<F: CstFacade> CstVisitor<F> for Unformatter {
type Error = Infallible;
fn visit_str(&mut self, handle: StrHandle, view: StrView, tree: &F) -> Result<(), Self::Error> {
self.no_white_space_here = true;
self.visit_str_super(handle, view, tree)?;
self.no_white_space_here = false;
Ok(())
}
fn visit_code_block(
&mut self,
handle: CodeBlockHandle,
view: CodeBlockView,
tree: &F,
) -> Result<(), Self::Error> {
self.no_white_space_here = true;
self.visit_code_block_super(handle, view, tree)?;
self.no_white_space_here = false;
Ok(())
}
fn visit_non_terminal(
&mut self,
id: CstNodeId,
_kind: NonTerminalKind,
_data: NonTerminalData,
_tree: &F,
) -> Result<(), Self::Error> {
self.parent = Some(id);
Ok(())
}
fn visit_terminal(
&mut self,
id: CstNodeId,
kind: TerminalKind,
_data: TerminalData,
_tree: &F,
) -> Result<(), Self::Error> {
if kind.is_builtin_terminal() {
return Ok(());
}
let Some(parent) = self.parent else {
return Ok(());
};
if !self.no_white_space_here && self.rng.random::<f32>() < self.weird_space_probability {
let spaces = (0..self.rng.random_range(1..3))
.map(|_| {
let ws = self.random_whitespace();
self.commands
.insert_dynamic_terminal(TerminalKind::Whitespace, ws)
})
.collect::<Vec<_>>();
self.commands.add_nodes_before(parent, id, spaces);
}
if self.rng.random::<f32>() < self.empty_line_probability {
let inserted = self
.commands
.insert_dynamic_terminal(TerminalKind::NewLine, "\n");
self.commands.add_nodes_before(parent, id, vec![inserted]);
}
Ok(())
}
fn visit_whitespace_terminal(
&mut self,
terminal: Whitespace,
_data: TerminalData,
_tree: &F,
) -> Result<(), Self::Error> {
if self.rng.random::<f32>() < self.whitespace_removal_probability {
self.commands.delete_node(terminal.node_id());
}
Ok(())
}
fn visit_new_line_terminal(
&mut self,
terminal: NewLine,
_data: TerminalData,
_tree: &F,
) -> Result<(), Self::Error> {
if self.rng.random::<f32>() < self.line_removal_probability {
self.commands.delete_node(terminal.node_id());
}
Ok(())
}
}