1use crate::nodes::{
5 Expression, FieldExpression, FunctionCall, IndexExpression, NumberExpression, Prefix,
6 Statement, StringSegment, TableExpression, Variable,
7};
8
9const QUOTED_STRING_MAX_LENGTH: usize = 60;
10const LONG_STRING_MIN_LENGTH: usize = 20;
11const FORCE_LONG_STRING_NEW_LINE_THRESHOLD: usize = 6;
12
13#[inline]
14pub fn should_break_with_space(ending_character: char, next_character: char) -> bool {
15 match ending_character {
16 '0'..='9' => matches!(next_character, '0'..='9' | 'A'..='Z' | 'a'..='z' | '_' | '.'),
17 'A'..='Z' | 'a'..='z' | '_' => {
18 next_character.is_ascii_alphanumeric() || next_character == '_'
19 }
20 '>' => next_character == '=',
21 '-' => next_character == '-',
22 '[' => next_character == '[',
23 ']' => next_character == ']',
24 '.' => matches!(next_character, '.' | '0'..='9'),
25 _ => false,
26 }
27}
28
29pub fn break_long_string(last_str: &str) -> bool {
30 if let Some(last_char) = last_str.chars().last() {
31 last_char == '['
32 } else {
33 false
34 }
35}
36
37pub fn break_variable_arguments(last_string: &str) -> bool {
38 if let Some('.') = last_string.chars().last() {
39 true
40 } else if let Some(first_char) = last_string.chars().next() {
41 first_char == '.' || first_char.is_ascii_digit()
42 } else {
43 false
44 }
45}
46
47pub fn break_minus(last_string: &str) -> bool {
48 if let Some(last_char) = last_string.chars().last() {
49 last_char == '-'
50 } else {
51 false
52 }
53}
54
55pub fn break_equal(last_string: &str) -> bool {
56 if let Some(last_char) = last_string.chars().last() {
57 last_char == '>'
58 } else {
59 false
60 }
61}
62
63pub fn break_concat(last_string: &str) -> bool {
64 if let Some('.') = last_string.chars().last() {
65 true
66 } else if let Some(first_char) = last_string.chars().next() {
67 first_char == '.' || first_char.is_ascii_digit()
68 } else {
69 false
70 }
71}
72
73pub fn ends_with_prefix(statement: &Statement) -> bool {
74 match statement {
75 Statement::Assign(assign) => {
76 if let Some(value) = assign.last_value() {
77 expression_ends_with_prefix(value)
78 } else {
79 false
80 }
81 }
82 Statement::CompoundAssign(assign) => expression_ends_with_prefix(assign.get_value()),
83 Statement::Call(_) => true,
84 Statement::Repeat(repeat) => expression_ends_with_prefix(repeat.get_condition()),
85 Statement::LocalAssign(assign) => {
86 if let Some(value) = assign.last_value() {
87 expression_ends_with_prefix(value)
88 } else {
89 false
90 }
91 }
92 _ => false,
93 }
94}
95
96pub fn starts_with_table(mut expression: &Expression) -> Option<&TableExpression> {
97 loop {
98 match expression {
99 Expression::Table(table) => break Some(table),
100 Expression::Binary(binary) => {
101 expression = binary.left();
102 }
103 Expression::Call(_)
104 | Expression::False(_)
105 | Expression::Field(_)
106 | Expression::Function(_)
107 | Expression::Identifier(_)
108 | Expression::If(_)
109 | Expression::Index(_)
110 | Expression::Nil(_)
111 | Expression::Number(_)
112 | Expression::Parenthese(_)
113 | Expression::String(_)
114 | Expression::InterpolatedString(_)
115 | Expression::True(_)
116 | Expression::Unary(_)
117 | Expression::VariableArguments(_) => break None,
118 Expression::TypeCast(type_cast) => {
119 expression = type_cast.get_expression();
120 }
121 }
122 }
123}
124
125pub fn starts_with_parenthese(statement: &Statement) -> bool {
126 match statement {
127 Statement::Assign(assign) => {
128 if let Some(variable) = assign.get_variables().first() {
129 match variable {
130 Variable::Identifier(_) => false,
131 Variable::Field(field) => field_starts_with_parenthese(field),
132 Variable::Index(index) => index_starts_with_parenthese(index),
133 }
134 } else {
135 false
136 }
137 }
138 Statement::CompoundAssign(assign) => match assign.get_variable() {
139 Variable::Identifier(_) => false,
140 Variable::Field(field) => field_starts_with_parenthese(field),
141 Variable::Index(index) => index_starts_with_parenthese(index),
142 },
143 Statement::Call(call) => call_starts_with_parenthese(call),
144 _ => false,
145 }
146}
147
148fn expression_ends_with_prefix(expression: &Expression) -> bool {
149 match expression {
150 Expression::Binary(binary) => expression_ends_with_prefix(binary.right()),
151 Expression::Call(_)
152 | Expression::Parenthese(_)
153 | Expression::Identifier(_)
154 | Expression::Field(_)
155 | Expression::Index(_) => true,
156 Expression::Unary(unary) => expression_ends_with_prefix(unary.get_expression()),
157 Expression::If(if_expression) => {
158 expression_ends_with_prefix(if_expression.get_else_result())
159 }
160 Expression::False(_)
161 | Expression::Function(_)
162 | Expression::Nil(_)
163 | Expression::Number(_)
164 | Expression::String(_)
165 | Expression::InterpolatedString(_)
166 | Expression::Table(_)
167 | Expression::True(_)
168 | Expression::VariableArguments(_)
169 | Expression::TypeCast(_) => false,
170 }
171}
172
173fn prefix_starts_with_parenthese(prefix: &Prefix) -> bool {
174 match prefix {
175 Prefix::Parenthese(_) => true,
176 Prefix::Call(call) => call_starts_with_parenthese(call),
177 Prefix::Field(field) => field_starts_with_parenthese(field),
178 Prefix::Index(index) => index_starts_with_parenthese(index),
179 Prefix::Identifier(_) => false,
180 }
181}
182
183#[inline]
184fn call_starts_with_parenthese(call: &FunctionCall) -> bool {
185 prefix_starts_with_parenthese(call.get_prefix())
186}
187
188#[inline]
189fn field_starts_with_parenthese(field: &FieldExpression) -> bool {
190 prefix_starts_with_parenthese(field.get_prefix())
191}
192
193#[inline]
194fn index_starts_with_parenthese(index: &IndexExpression) -> bool {
195 prefix_starts_with_parenthese(index.get_prefix())
196}
197
198pub fn write_number(number: &NumberExpression) -> String {
199 match number {
200 NumberExpression::Decimal(number) => {
201 let float = number.get_raw_float();
202 if float.is_nan() {
203 "(0/0)".to_owned()
204 } else if float.is_infinite() {
205 format!("({}1/0)", if float.is_sign_negative() { "-" } else { "" })
206 } else {
207 format!(
208 "{}{}",
209 float,
210 number
211 .get_exponent()
212 .map(|exponent| {
213 let exponent_char = number
214 .is_uppercase()
215 .map(|is_uppercase| if is_uppercase { 'E' } else { 'e' })
216 .unwrap_or('e');
217 format!("{}{}", exponent_char, exponent)
218 })
219 .unwrap_or_else(|| "".to_owned())
220 )
221 }
222 }
223 NumberExpression::Hex(number) => {
224 format!(
225 "0{}{:x}{}",
226 if number.is_x_uppercase() { 'X' } else { 'x' },
227 number.get_raw_integer(),
228 number
229 .get_exponent()
230 .map(|exponent| {
231 let exponent_char = number
232 .is_exponent_uppercase()
233 .map(|is_uppercase| if is_uppercase { 'P' } else { 'p' })
234 .unwrap_or('p');
235 format!("{}{}", exponent_char, exponent)
236 })
237 .unwrap_or_else(|| "".to_owned())
238 )
239 }
240 NumberExpression::Binary(number) => {
241 format!(
242 "0{}{:b}",
243 if number.is_b_uppercase() { 'B' } else { 'b' },
244 number.get_raw_value()
245 )
246 }
247 }
248}
249
250fn needs_escaping(character: char) -> bool {
251 !(character.is_ascii_graphic() || character == ' ') || character == '\\'
252}
253
254fn needs_quoted_string(character: char) -> bool {
255 !(character.is_ascii_graphic() || character == ' ' || character == '\n')
256}
257
258fn escape(character: char) -> String {
259 match character {
260 '\n' => "\\n".to_owned(),
261 '\t' => "\\t".to_owned(),
262 '\\' => "\\\\".to_owned(),
263 '\r' => "\\r".to_owned(),
264 '\u{7}' => "\\a".to_owned(),
265 '\u{8}' => "\\b".to_owned(),
266 '\u{B}' => "\\v".to_owned(),
267 '\u{C}' => "\\f".to_owned(),
268 _ => {
269 if (character as u32) < 256 {
270 format!("\\{}", character as u8)
271 } else {
272 format!("\\u{{{:x}}}", character as u32)
273 }
274 }
275 }
276}
277
278#[inline]
279pub fn count_new_lines(string: &str) -> usize {
280 string.chars().filter(|c| *c == '\n').count()
281}
282
283pub fn write_string(value: &str) -> String {
284 if value.is_empty() {
285 return "''".to_owned();
286 }
287
288 if value.len() == 1 {
289 let character = value
290 .chars()
291 .next()
292 .expect("string should have at least one character");
293 match character {
294 '\'' => return "\"'\"".to_owned(),
295 '"' => return "'\"'".to_owned(),
296 _ => {
297 if needs_escaping(character) {
298 return format!("'{}'", escape(character));
299 } else {
300 return format!("'{}'", character);
301 }
302 }
303 }
304 }
305
306 if !value.contains(needs_quoted_string)
307 && value.len() >= LONG_STRING_MIN_LENGTH
308 && (value.len() >= QUOTED_STRING_MAX_LENGTH
309 || count_new_lines(value) >= FORCE_LONG_STRING_NEW_LINE_THRESHOLD)
310 {
311 write_long_bracket(value)
312 } else {
313 write_quoted(value)
314 }
315}
316
317pub fn write_interpolated_string_segment(segment: &StringSegment) -> String {
318 let value = segment.get_value();
319
320 if value.is_empty() {
321 return "".to_owned();
322 }
323
324 let mut result = String::new();
325
326 result.reserve(value.len());
327
328 for character in value.chars() {
329 match character {
330 '`' | '{' => {
331 result.push('\\');
332 result.push(character);
333 }
334 _ if needs_escaping(character) => {
335 result.push_str(&escape(character));
336 }
337 _ => {
338 result.push(character);
339 }
340 }
341 }
342
343 result
344}
345
346fn write_long_bracket(value: &str) -> String {
347 let mut i: usize = value.ends_with(']').into();
348 let mut equals = "=".repeat(i);
349 loop {
350 if !value.contains(&format!("]{}]", equals)) {
351 break;
352 } else {
353 i += 1;
354 equals = "=".repeat(i);
355 };
356 }
357 let needs_extra_new_line = if value.starts_with('\n') { "\n" } else { "" };
358 format!("[{}[{}{}]{}]", equals, needs_extra_new_line, value, equals)
359}
360
361fn write_quoted(value: &str) -> String {
362 let mut quoted = String::new();
363 quoted.reserve(value.len() + 2);
364
365 let quote_symbol = get_quote_symbol(value);
366 quoted.push(quote_symbol);
367
368 for character in value.chars() {
369 if character == quote_symbol {
370 quoted.push('\\');
371 quoted.push(quote_symbol);
372 } else if needs_escaping(character) {
373 quoted.push_str(&escape(character));
374 } else {
375 quoted.push(character);
376 }
377 }
378
379 quoted.push(quote_symbol);
380 quoted.shrink_to_fit();
381 quoted
382}
383
384fn get_quote_symbol(value: &str) -> char {
385 if value.contains('"') {
386 '\''
387 } else if value.contains('\'') {
388 '"'
389 } else {
390 '\''
391 }
392}
393
394#[cfg(test)]
395mod test {
396 use super::*;
397
398 mod write_string {
399 use super::*;
400
401 macro_rules! test_output {
402 ($($name:ident($input:literal) => $value:literal),* $(,)?) => {
403 $(
404 #[test]
405 fn $name() {
406 assert_eq!($value, write_string(&$input));
407 }
408 )*
409 };
410 }
411
412 test_output!(
413 empty("") => "''",
414 single_letter("a") => "'a'",
415 single_digit("8") => "'8'",
416 single_symbol("!") => "'!'",
417 single_space(" ") => "' '",
418 abc("abc") => "'abc'",
419 three_spaces(" ") => "' '",
420 new_line("\n") => "'\\n'",
421 bell("\u{7}") => "'\\a'",
422 backspace("\u{8}") => "'\\b'",
423 form_feed("\u{c}") => "'\\f'",
424 tab("\t") => "'\\t'",
425 carriage_return("\u{D}") => "'\\r'",
426 vertical_tab("\u{B}") => "'\\v'",
427 backslash("\\") => "'\\\\'",
428 single_quote("'") => "\"'\"",
429 double_quote("\"") => "'\"'",
430 null("\0") => "'\\0'",
431 escape("\u{1B}") => "'\\27'",
432 extended_ascii("\u{C3}") => "'\\195'",
433 unicode("\u{25C1}") => "'\\u{25c1}'",
434 im_cool("I'm cool") => "\"I'm cool\"",
435 ends_with_closing_bracket("oof]") => "'oof]'",
436 multiline_ends_with_closing_bracket("oof\noof]") => "'oof\\noof]'",
437 large_multiline_does_not_end_with_closing_bracket("ooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof")
438 => "[[ooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof]]",
439 large_multiline_ends_with_closing_bracket("ooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof]")
440 => "[=[ooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof]]=]",
441 large_multiline_starts_with_new_line("\nooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof")
442 => "[[\n\nooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof]]",
443
444 large_multiline_with_unicode("\nooof\nooof\nooof\nooof\nooof\nooof\nooof\nooof\noof\u{10FFFF}")
445 => "'\\nooof\\nooof\\nooof\\nooof\\nooof\\nooof\\nooof\\nooof\\noof\\u{10ffff}'",
446 );
447 }
448}