feat: generalize removing and adding back prefix

src/formatting.rs

@@ -3,20 +3,19 @@ use rug::Float;
 use crate::nodes::{self, Error};
 
 pub fn pretty_eval(input: &str, max_digits: usize) -> Result<String, Error> {
-    let mut input = input.replace('_', "");
+    let (prefix, input) = input
+        .rsplit_once('=')
+        .map(|(p, i)| (p.to_owned() + "= ", i))
+        .unwrap_or_else(|| {
+            input
+                .rsplit_once('&')
+                .map(|(p, i)| (p.to_owned() + "& ", i))
+                .unwrap_or((String::new(), input))
+        });
 
-    // no need to remove the characters here,
-    // as we ignore the them in all other logic already
-    let should_add_back_alignment = input.contains('&');
+    // no need to remove the character here,
+    // as we ignore it in all other logic already
     let should_add_back_backslash = input.contains('\\');
-    let leading_spaces = input.chars().take_while(|c| *c == ' ').count();
-
-    let should_add_back_equals = input
-        .find('=')
-        .inspect(|i| {
-            input.remove(*i);
-        })
-        .is_some();
 
     let result = nodes::eval(&input)?;
     let display_digits =
@@ -34,17 +33,7 @@ pub fn pretty_eval(input: &str, max_digits: usize) -> Result<String, Error> {
             });
 
     Ok(format!(
-        "{}{}{result:.display_digits$}{}",
-        " ".repeat(leading_spaces),
-        if should_add_back_equals && should_add_back_alignment {
-            "& = "
-        } else if should_add_back_equals {
-            "= "
-        } else if should_add_back_alignment {
-            "& "
-        } else {
-            ""
-        },
+        "{prefix}{result:.display_digits$}{}",
         if should_add_back_backslash { " \\" } else { "" }
     )
     .replace('e', " * 10^"))