From acc49edd3178b87db3acd36b39aca5dafd110d93 Mon Sep 17 00:00:00 2001
From: electria <elec3137+git@proton.me>
Date: Fri, 9 Oct 2026 18:50:19 -0700
Subject: [PATCH] feat: generalize removing and adding back prefix

---
 src/formatting.rs | 37 +++++++++++++------------------------
 1 file changed, 13 insertions(+), 24 deletions(-)

diff --git a/src/formatting.rs b/src/formatting.rs
index dcb740f..fff48dc 100644
--- a/src/formatting.rs
+++ b/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^"))
@@ -93,7 +82,7 @@ mod tests {
     #[test]
     fn character_readdition() {
         assert_eq!(pretty_eval(" = 11", 10).unwrap(), " = 11");
-        assert_eq!(pretty_eval(" &= 11", 10).unwrap(), " & = 11");
+        assert_eq!(pretty_eval(" &= 11", 10).unwrap(), " &= 11");
         assert_eq!(pretty_eval(" & 11", 10).unwrap(), " & 11");
         assert_eq!(pretty_eval(" & 11 \\", 10).unwrap(), " & 11 \\");
     }