From 8509f15427d09354f4fe11188b9b5b6fc4329840 Mon Sep 17 00:00:00 2001
From: Robert Bennett <rltbennett@icloud.com>
Date: Mon, 28 Oct 2024 22:06:33 -0400
Subject: [PATCH] Clean up uses of `CompactString` Reduce/removed some
 allocations

---
 numbat/examples/inspect.rs |  10 +--
 numbat/src/ast.rs          |   5 +-
 numbat/src/datetime.rs     |  27 +++---
 numbat/src/number.rs       |   5 +-
 numbat/src/parser.rs       |  11 ++-
 numbat/src/prefix.rs       | 164 ++++++++++++++++++-------------------
 numbat/src/pretty_print.rs |  30 ++++---
 numbat/src/resolver.rs     |  15 ++--
 8 files changed, 143 insertions(+), 124 deletions(-)

diff --git a/numbat/examples/inspect.rs b/numbat/examples/inspect.rs
index 2efda6ba..3c8c9d15 100644
--- a/numbat/examples/inspect.rs
+++ b/numbat/examples/inspect.rs
@@ -1,4 +1,4 @@
-use compact_str::{format_compact, CompactString, ToCompactString};
+use compact_str::{format_compact, CompactString};
 use itertools::Itertools;
 use numbat::markup::plain_text_format;
 use numbat::module_importer::FileSystemImporter;
@@ -61,8 +61,7 @@ fn inspect_functions_in_module(ctx: &Context, prelude_ctx: &Context, module: Str
         }
 
         if let Some(ref description_raw) = description {
-            let description =
-                replace_equation_delimiters(description_raw.trim().to_compact_string());
+            let description = replace_equation_delimiters(description_raw.trim());
 
             if description.ends_with('.') {
                 println!("{description}");
@@ -128,7 +127,7 @@ fn inspect_functions_in_module(ctx: &Context, prelude_ctx: &Context, module: Str
 
                     //Print the example
                     if let Some(example_description) = example_description {
-                        println!("{}", replace_equation_delimiters(example_description));
+                        println!("{}", replace_equation_delimiters(&example_description));
                     }
 
                     print!("<pre>");
@@ -159,8 +158,9 @@ fn inspect_functions_in_module(ctx: &Context, prelude_ctx: &Context, module: Str
 }
 
 // Replace $..$ with \\( .. \\) for mdbook.
-fn replace_equation_delimiters(text_in: CompactString) -> CompactString {
+fn replace_equation_delimiters(text_in: &str) -> CompactString {
     let mut text_out = CompactString::with_capacity(text_in.len());
+    // TODO: handle \$ in math
     for (i, part) in text_in.split('$').enumerate() {
         if i % 2 == 0 {
             text_out.push_str(part);
diff --git a/numbat/src/ast.rs b/numbat/src/ast.rs
index 18bf6114..3f476063 100644
--- a/numbat/src/ast.rs
+++ b/numbat/src/ast.rs
@@ -1,8 +1,9 @@
 use crate::markup as m;
+use crate::resolver::ModulePathBorrowed;
 use crate::span::Span;
 use crate::{
     arithmetic::Exponent, decorator::Decorator, markup::Markup, number::Number, prefix::Prefix,
-    pretty_print::PrettyPrint, resolver::ModulePath,
+    pretty_print::PrettyPrint,
 };
 use compact_str::{format_compact, CompactString, ToCompactString};
 use itertools::Itertools;
@@ -444,7 +445,7 @@ pub enum Statement<'a> {
         decorators: Vec<Decorator<'a>>,
     },
     ProcedureCall(Span, ProcedureKind, Vec<Expression<'a>>),
-    ModuleImport(Span, ModulePath),
+    ModuleImport(Span, ModulePathBorrowed<'a>),
     DefineStruct {
         struct_name_span: Span,
         struct_name: &'a str,
diff --git a/numbat/src/datetime.rs b/numbat/src/datetime.rs
index 43b0fdc0..5f3da415 100644
--- a/numbat/src/datetime.rs
+++ b/numbat/src/datetime.rs
@@ -1,4 +1,4 @@
-use compact_str::{format_compact, CompactString, ToCompactString};
+use compact_str::{CompactString, ToCompactString};
 use jiff::{civil::DateTime, fmt::rfc2822, tz::TimeZone, Timestamp, Zoned};
 use std::str::FromStr;
 
@@ -75,22 +75,23 @@ pub fn to_string(dt: &Zoned) -> CompactString {
     if dt.time_zone() == &TimeZone::UTC {
         dt.strftime("%Y-%m-%d %H:%M:%S UTC").to_compact_string()
     } else {
+        use std::fmt::Write;
+        let mut out = CompactString::with_capacity("2000-01-01 00:00:00 (UTC +00:00)".len());
+        write!(out, "{}", dt.strftime("%Y-%m-%d %H:%M:%S")).unwrap();
+
         let offset = dt.offset();
         let zone_abbreviation = tz.to_offset(dt.timestamp()).2;
-        let abbreviation_and_offset =
-            if zone_abbreviation.starts_with('+') || zone_abbreviation.starts_with('-') {
-                format_compact!("(UTC {offset})")
-            } else {
-                format_compact!("{zone_abbreviation} (UTC {offset})")
-            };
-
-        let timezone_name = if let Some(iana_tz_name) = tz.iana_name() {
-            format_compact!(", {iana_tz_name}")
+
+        if zone_abbreviation.starts_with('+') || zone_abbreviation.starts_with('-') {
+            write!(out, " (UTC {offset})").unwrap();
         } else {
-            CompactString::const_new("")
+            write!(out, " {zone_abbreviation} (UTC {offset})").unwrap();
         };
 
-        let dt_str = dt.strftime("%Y-%m-%d %H:%M:%S");
-        format_compact!("{dt_str} {abbreviation_and_offset}{timezone_name}")
+        if let Some(iana_tz_name) = tz.iana_name() {
+            write!(out, ", {iana_tz_name}").unwrap();
+        }
+
+        out
     }
 }
diff --git a/numbat/src/number.rs b/numbat/src/number.rs
index 8e119fca..47286396 100644
--- a/numbat/src/number.rs
+++ b/numbat/src/number.rs
@@ -83,8 +83,7 @@ impl Number {
                     .round()
             };
 
-            // TODO: wasteful to go through a String here
-            let formatted_number = dtoa(number, config).to_compact_string();
+            let formatted_number = dtoa(number, config);
 
             if formatted_number.contains('.') && !formatted_number.contains('e') {
                 let formatted_number = if config.max_sig_digits.is_some() {
@@ -101,7 +100,7 @@ impl Number {
             } else if formatted_number.contains('e') && !formatted_number.contains("e-") {
                 formatted_number.replace('e', "e+").to_compact_string()
             } else {
-                formatted_number
+                formatted_number.to_compact_string()
             }
         }
     }
diff --git a/numbat/src/parser.rs b/numbat/src/parser.rs
index 05645768..4204f60e 100644
--- a/numbat/src/parser.rs
+++ b/numbat/src/parser.rs
@@ -70,7 +70,7 @@ use crate::ast::{
 use crate::decorator::{self, Decorator};
 use crate::number::Number;
 use crate::prefix_parser::AcceptsPrefix;
-use crate::resolver::ModulePath;
+use crate::resolver::ModulePathBorrowed;
 use crate::span::Span;
 use crate::tokenizer::{Token, TokenKind, TokenizerError, TokenizerErrorKind};
 
@@ -898,11 +898,11 @@ impl<'a> Parser<'a> {
         let mut span = self.peek(tokens).span;
 
         if let Some(identifier) = self.match_exact(tokens, TokenKind::Identifier) {
-            let mut module_path = vec![identifier.lexeme.to_compact_string()];
+            let mut module_path = vec![identifier.lexeme];
 
             while self.match_exact(tokens, TokenKind::DoubleColon).is_some() {
                 if let Some(identifier) = self.match_exact(tokens, TokenKind::Identifier) {
-                    module_path.push(identifier.lexeme.to_compact_string());
+                    module_path.push(identifier.lexeme);
                 } else {
                     return Err(ParseError {
                         kind: ParseErrorKind::ExpectedModuleNameAfterDoubleColon,
@@ -912,7 +912,10 @@ impl<'a> Parser<'a> {
             }
             span = span.extend(&self.last(tokens).unwrap().span);
 
-            Ok(Statement::ModuleImport(span, ModulePath(module_path)))
+            Ok(Statement::ModuleImport(
+                span,
+                ModulePathBorrowed(module_path),
+            ))
         } else {
             Err(ParseError {
                 kind: ParseErrorKind::ExpectedModulePathAfterUse,
diff --git a/numbat/src/prefix.rs b/numbat/src/prefix.rs
index 53020d39..76e37538 100644
--- a/numbat/src/prefix.rs
+++ b/numbat/src/prefix.rs
@@ -105,90 +105,90 @@ impl Prefix {
     }
 
     pub fn as_string_short(&self) -> CompactString {
-        match self {
-            Prefix::Metric(-30) => CompactString::const_new("q"),
-            Prefix::Metric(-27) => CompactString::const_new("r"),
-            Prefix::Metric(-24) => CompactString::const_new("y"),
-            Prefix::Metric(-21) => CompactString::const_new("z"),
-            Prefix::Metric(-18) => CompactString::const_new("a"),
-            Prefix::Metric(-15) => CompactString::const_new("f"),
-            Prefix::Metric(-12) => CompactString::const_new("p"),
-            Prefix::Metric(-9) => CompactString::const_new("n"),
-            Prefix::Metric(-6) => CompactString::const_new("µ"),
-            Prefix::Metric(-3) => CompactString::const_new("m"),
-            Prefix::Metric(-2) => CompactString::const_new("c"),
-            Prefix::Metric(-1) => CompactString::const_new("d"),
-            Prefix::Metric(0) => CompactString::const_new(""),
-            Prefix::Metric(1) => CompactString::const_new("da"),
-            Prefix::Metric(2) => CompactString::const_new("h"),
-            Prefix::Metric(3) => CompactString::const_new("k"),
-            Prefix::Metric(6) => CompactString::const_new("M"),
-            Prefix::Metric(9) => CompactString::const_new("G"),
-            Prefix::Metric(12) => CompactString::const_new("T"),
-            Prefix::Metric(15) => CompactString::const_new("P"),
-            Prefix::Metric(18) => CompactString::const_new("E"),
-            Prefix::Metric(21) => CompactString::const_new("Z"),
-            Prefix::Metric(24) => CompactString::const_new("Y"),
-            Prefix::Metric(27) => CompactString::const_new("R"),
-            Prefix::Metric(30) => CompactString::const_new("Q"),
-
-            Prefix::Metric(n) => format_compact!("<prefix 10^{n}>"),
-
-            Prefix::Binary(0) => CompactString::const_new(""),
-            Prefix::Binary(10) => CompactString::const_new("Ki"),
-            Prefix::Binary(20) => CompactString::const_new("Mi"),
-            Prefix::Binary(30) => CompactString::const_new("Gi"),
-            Prefix::Binary(40) => CompactString::const_new("Ti"),
-            Prefix::Binary(50) => CompactString::const_new("Pi"),
-            Prefix::Binary(60) => CompactString::const_new("Ei"),
-            Prefix::Binary(70) => CompactString::const_new("Zi"),
-            Prefix::Binary(80) => CompactString::const_new("Yi"),
-
-            Prefix::Binary(n) => format_compact!("<prefix 2^{n}>"),
-        }
+        CompactString::const_new(match self {
+            Prefix::Metric(-30) => "q",
+            Prefix::Metric(-27) => "r",
+            Prefix::Metric(-24) => "y",
+            Prefix::Metric(-21) => "z",
+            Prefix::Metric(-18) => "a",
+            Prefix::Metric(-15) => "f",
+            Prefix::Metric(-12) => "p",
+            Prefix::Metric(-9) => "n",
+            Prefix::Metric(-6) => "µ",
+            Prefix::Metric(-3) => "m",
+            Prefix::Metric(-2) => "c",
+            Prefix::Metric(-1) => "d",
+            Prefix::Metric(0) => "",
+            Prefix::Metric(1) => "da",
+            Prefix::Metric(2) => "h",
+            Prefix::Metric(3) => "k",
+            Prefix::Metric(6) => "M",
+            Prefix::Metric(9) => "G",
+            Prefix::Metric(12) => "T",
+            Prefix::Metric(15) => "P",
+            Prefix::Metric(18) => "E",
+            Prefix::Metric(21) => "Z",
+            Prefix::Metric(24) => "Y",
+            Prefix::Metric(27) => "R",
+            Prefix::Metric(30) => "Q",
+
+            Prefix::Metric(n) => return format_compact!("<prefix 10^{n}>"),
+
+            Prefix::Binary(0) => "",
+            Prefix::Binary(10) => "Ki",
+            Prefix::Binary(20) => "Mi",
+            Prefix::Binary(30) => "Gi",
+            Prefix::Binary(40) => "Ti",
+            Prefix::Binary(50) => "Pi",
+            Prefix::Binary(60) => "Ei",
+            Prefix::Binary(70) => "Zi",
+            Prefix::Binary(80) => "Yi",
+
+            Prefix::Binary(n) => return format_compact!("<prefix 2^{n}>"),
+        })
     }
 
     pub fn as_string_long(&self) -> CompactString {
-        match self {
-            Prefix::Metric(-30) => CompactString::const_new("quecto"),
-            Prefix::Metric(-27) => CompactString::const_new("ronto"),
-            Prefix::Metric(-24) => CompactString::const_new("yocto"),
-            Prefix::Metric(-21) => CompactString::const_new("zepto"),
-            Prefix::Metric(-18) => CompactString::const_new("atto"),
-            Prefix::Metric(-15) => CompactString::const_new("femto"),
-            Prefix::Metric(-12) => CompactString::const_new("pico"),
-            Prefix::Metric(-9) => CompactString::const_new("nano"),
-            Prefix::Metric(-6) => CompactString::const_new("micro"),
-            Prefix::Metric(-3) => CompactString::const_new("milli"),
-            Prefix::Metric(-2) => CompactString::const_new("centi"),
-            Prefix::Metric(-1) => CompactString::const_new("deci"),
-            Prefix::Metric(0) => CompactString::const_new(""),
-            Prefix::Metric(1) => CompactString::const_new("deca"),
-            Prefix::Metric(2) => CompactString::const_new("hecto"),
-            Prefix::Metric(3) => CompactString::const_new("kilo"),
-            Prefix::Metric(6) => CompactString::const_new("mega"),
-            Prefix::Metric(9) => CompactString::const_new("giga"),
-            Prefix::Metric(12) => CompactString::const_new("tera"),
-            Prefix::Metric(15) => CompactString::const_new("peta"),
-            Prefix::Metric(18) => CompactString::const_new("exa"),
-            Prefix::Metric(21) => CompactString::const_new("zetta"),
-            Prefix::Metric(24) => CompactString::const_new("yotta"),
-            Prefix::Metric(27) => CompactString::const_new("ronna"),
-            Prefix::Metric(30) => CompactString::const_new("quetta"),
-
-            Prefix::Metric(n) => format_compact!("<prefix 10^{n}>"),
-
-            Prefix::Binary(0) => CompactString::const_new(""),
-            Prefix::Binary(10) => CompactString::const_new("kibi"),
-            Prefix::Binary(20) => CompactString::const_new("mebi"),
-            Prefix::Binary(30) => CompactString::const_new("gibi"),
-            Prefix::Binary(40) => CompactString::const_new("tebi"),
-            Prefix::Binary(50) => CompactString::const_new("pebi"),
-            Prefix::Binary(60) => CompactString::const_new("exbi"),
-            Prefix::Binary(70) => CompactString::const_new("zebi"),
-            Prefix::Binary(80) => CompactString::const_new("yobi"),
-
-            Prefix::Binary(n) => format_compact!("<prefix 2^{n}>"),
-        }
+        CompactString::const_new(match self {
+            Prefix::Metric(-30) => "quecto",
+            Prefix::Metric(-27) => "ronto",
+            Prefix::Metric(-24) => "yocto",
+            Prefix::Metric(-21) => "zepto",
+            Prefix::Metric(-18) => "atto",
+            Prefix::Metric(-15) => "femto",
+            Prefix::Metric(-12) => "pico",
+            Prefix::Metric(-9) => "nano",
+            Prefix::Metric(-6) => "micro",
+            Prefix::Metric(-3) => "milli",
+            Prefix::Metric(-2) => "centi",
+            Prefix::Metric(-1) => "deci",
+            Prefix::Metric(0) => "",
+            Prefix::Metric(1) => "deca",
+            Prefix::Metric(2) => "hecto",
+            Prefix::Metric(3) => "kilo",
+            Prefix::Metric(6) => "mega",
+            Prefix::Metric(9) => "giga",
+            Prefix::Metric(12) => "tera",
+            Prefix::Metric(15) => "peta",
+            Prefix::Metric(18) => "exa",
+            Prefix::Metric(21) => "zetta",
+            Prefix::Metric(24) => "yotta",
+            Prefix::Metric(27) => "ronna",
+            Prefix::Metric(30) => "quetta",
+
+            Prefix::Metric(n) => return format_compact!("<prefix 10^{n}>"),
+
+            Prefix::Binary(0) => "",
+            Prefix::Binary(10) => "kibi",
+            Prefix::Binary(20) => "mebi",
+            Prefix::Binary(30) => "gibi",
+            Prefix::Binary(40) => "tebi",
+            Prefix::Binary(50) => "pebi",
+            Prefix::Binary(60) => "exbi",
+            Prefix::Binary(70) => "zebi",
+            Prefix::Binary(80) => "yobi",
+
+            Prefix::Binary(n) => return format_compact!("<prefix 2^{n}>"),
+        })
     }
 }
diff --git a/numbat/src/pretty_print.rs b/numbat/src/pretty_print.rs
index 113c0c02..556893b0 100644
--- a/numbat/src/pretty_print.rs
+++ b/numbat/src/pretty_print.rs
@@ -1,4 +1,4 @@
-use compact_str::{CompactString, ToCompactString};
+use compact_str::CompactString;
 
 use crate::markup::Markup;
 
@@ -14,15 +14,25 @@ impl PrettyPrint for bool {
 }
 
 pub fn escape_numbat_string(s: &str) -> CompactString {
-    s.replace("\\", "\\\\")
-        .replace("\n", "\\n")
-        .replace("\r", "\\r")
-        .replace("\t", "\\t")
-        .replace("\"", "\\\"")
-        .replace("\0", "\\0")
-        .replace("{", "\\{")
-        .replace("}", "\\}")
-        .to_compact_string()
+    let mut out = CompactString::const_new("");
+    for c in s.chars() {
+        let replacement = match c {
+            '\\' => r"\\",
+            '\n' => r"\n",
+            '\r' => r"\r",
+            '\t' => r"\t",
+            '"' => r#"\""#,
+            '\0' => r"\0",
+            '{' => r"\{",
+            '}' => r"\}",
+            _ => {
+                out.push(c);
+                continue;
+            }
+        };
+        out.push_str(replacement);
+    }
+    out
 }
 
 impl PrettyPrint for CompactString {
diff --git a/numbat/src/resolver.rs b/numbat/src/resolver.rs
index 987000c2..4af6426d 100644
--- a/numbat/src/resolver.rs
+++ b/numbat/src/resolver.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 
 use codespan_reporting::files::SimpleFiles;
-use compact_str::CompactString;
+use compact_str::{CompactString, ToCompactString};
 use thiserror::Error;
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -17,6 +17,9 @@ impl std::fmt::Display for ModulePath {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct ModulePathBorrowed<'a>(pub Vec<&'a str>);
+
 #[derive(Debug, Clone)]
 pub enum CodeSource {
     /// User input from the command line or a REPL
@@ -106,13 +109,15 @@ impl Resolver {
 
         for statement in program {
             match statement {
-                Statement::ModuleImport(span, module_path) => {
-                    if !self.imported_modules.contains(module_path) {
-                        if let Some((code, filesystem_path)) = self.importer.import(module_path) {
+                Statement::ModuleImport(span, ModulePathBorrowed(module_path)) => {
+                    if !self.imported_modules.iter().any(|m| &m.0 == module_path) {
+                        let module_path =
+                            ModulePath(module_path.iter().map(|s| s.to_compact_string()).collect());
+                        if let Some((code, filesystem_path)) = self.importer.import(&module_path) {
                             let code: &'static str = Box::leak(code.to_string().into_boxed_str());
                             self.imported_modules.push(module_path.clone());
                             let code_source_id = self.add_code_source(
-                                CodeSource::Module(module_path.clone(), filesystem_path),
+                                CodeSource::Module(module_path, filesystem_path),
                                 code,
                             );