Properly support recursive table formatting and multithreaded use

crates/lune-utils/src/fmt/value/config.rs

@@ -12,15 +12,18 @@ impl ValueFormatConfig {
         Creates a new config with default values.
     */
     #[must_use]
-    pub fn new() -> Self {
-        Self::default()
+    pub const fn new() -> Self {
+        Self {
+            max_depth: 3,
+            colors_enabled: false,
+        }
     }
 
     /**
         Sets the maximum depth to which tables will be formatted.
     */
     #[must_use]
-    pub fn with_max_depth(self, max_depth: usize) -> Self {
+    pub const fn with_max_depth(self, max_depth: usize) -> Self {
         Self { max_depth, ..self }
     }
 
@@ -30,7 +33,7 @@ impl ValueFormatConfig {
         Colors are disabled by default.
     */
     #[must_use]
-    pub fn with_colors_enabled(self, colors_enabled: bool) -> Self {
+    pub const fn with_colors_enabled(self, colors_enabled: bool) -> Self {
         Self {
             colors_enabled,
             ..self
@@ -40,9 +43,6 @@ impl ValueFormatConfig {
 
 impl Default for ValueFormatConfig {
     fn default() -> Self {
-        Self {
-            max_depth: 3,
-            colors_enabled: false,
-        }
+        Self::new()
     }
 }

crates/lune-utils/src/fmt/value/mod.rs

@@ -1,67 +1,26 @@
-use std::fmt::{self, Write as _};
+use std::{
+    collections::HashSet,
+    sync::{Arc, Mutex},
+};
 
-use console::{colors_enabled, set_colors_enabled};
+use console::{colors_enabled as get_colors_enabled, set_colors_enabled};
 use mlua::prelude::*;
+use once_cell::sync::Lazy;
 
 mod basic;
 mod config;
 mod metamethods;
+mod recursive;
 mod style;
 
+use self::recursive::format_value_recursive;
+
 pub use self::config::ValueFormatConfig;
 
-use self::basic::{format_value_styled, lua_value_as_plain_string_key};
-use self::style::STYLE_DIM;
-
-// NOTE: We return a result here but it's really just to make handling
-// of the `write!` calls easier. Writing into a string should never fail.
-fn format_value_inner(
-    value: &LuaValue,
-    config: &ValueFormatConfig,
-    depth: usize,
-) -> Result<String, fmt::Error> {
-    let mut buffer = String::new();
-
-    // TODO: Rewrite this section to not be recursive and
-    // keep track of any recursive references to tables.
-    if let LuaValue::Table(ref t) = value {
-        if depth >= config.max_depth {
-            write!(buffer, "{}", STYLE_DIM.apply_to("{ ... }"))?;
-        } else {
-            writeln!(buffer, "{}", STYLE_DIM.apply_to("{"))?;
-
-            for res in t.clone().pairs::<LuaValue, LuaValue>() {
-                let (key, value) = res.expect("conversion to LuaValue should never fail");
-                let formatted = if let Some(plain_key) = lua_value_as_plain_string_key(&key) {
-                    format!(
-                        "{} {} {}{}",
-                        plain_key,
-                        STYLE_DIM.apply_to("="),
-                        format_value_inner(&value, config, depth + 1)?,
-                        STYLE_DIM.apply_to(","),
-                    )
-                } else {
-                    format!(
-                        "{}{}{} {} {}{}",
-                        STYLE_DIM.apply_to("["),
-                        format_value_inner(&key, config, depth + 1)?,
-                        STYLE_DIM.apply_to("]"),
-                        STYLE_DIM.apply_to("="),
-                        format_value_inner(&value, config, depth + 1)?,
-                        STYLE_DIM.apply_to(","),
-                    )
-                };
-                buffer.push_str(&formatted);
-            }
-
-            writeln!(buffer, "{}", STYLE_DIM.apply_to("}"))?;
-        }
-    } else {
-        write!(buffer, "{}", format_value_styled(value))?;
-    }
-
-    Ok(buffer)
-}
+// NOTE: Since the setting for colors being enabled is global,
+// and these functions may be called in parallel, we use this global
+// lock to make sure that we don't mess up the colors for other threads.
+static COLORS_LOCK: Lazy<Arc<Mutex<()>>> = Lazy::new(|| Arc::new(Mutex::new(())));
 
 /**
     Formats a Lua value into a pretty string using the given config.
@@ -69,10 +28,15 @@ fn format_value_inner(
 #[must_use]
 #[allow(clippy::missing_panics_doc)]
 pub fn pretty_format_value(value: &LuaValue, config: &ValueFormatConfig) -> String {
-    let colors_were_enabled = colors_enabled();
-    set_colors_enabled(config.colors_enabled);
-    let res = format_value_inner(value, config, 0);
-    set_colors_enabled(colors_were_enabled);
+    let _guard = COLORS_LOCK.lock().unwrap();
+
+    let were_colors_enabled = get_colors_enabled();
+    set_colors_enabled(were_colors_enabled && config.colors_enabled);
+
+    let mut visited = HashSet::new();
+    let res = format_value_recursive(value, config, &mut visited, 0);
+
+    set_colors_enabled(were_colors_enabled);
     res.expect("using fmt for writing into strings should never fail")
 }
 
@@ -84,9 +48,18 @@ pub fn pretty_format_value(value: &LuaValue, config: &ValueFormatConfig) -> Stri
 #[must_use]
 #[allow(clippy::missing_panics_doc)]
 pub fn pretty_format_multi_value(values: &LuaMultiValue, config: &ValueFormatConfig) -> String {
-    values
+    let _guard = COLORS_LOCK.lock().unwrap();
+
+    let were_colors_enabled = get_colors_enabled();
+    set_colors_enabled(were_colors_enabled && config.colors_enabled);
+
+    let mut visited = HashSet::new();
+    let res = values
         .into_iter()
-        .map(|value| pretty_format_value(value, config))
-        .collect::<Vec<_>>()
+        .map(|value| format_value_recursive(value, config, &mut visited, 0))
+        .collect::<Result<Vec<_>, _>>();
+
+    set_colors_enabled(were_colors_enabled);
+    res.expect("using fmt for writing into strings should never fail")
         .join(" ")
 }

crates/lune-utils/src/fmt/value/recursive.rs

@@ -0,0 +1,84 @@
+use std::collections::HashSet;
+use std::fmt::{self, Write as _};
+
+use mlua::prelude::*;
+
+use super::{
+    basic::{format_value_styled, lua_value_as_plain_string_key},
+    config::ValueFormatConfig,
+    style::STYLE_DIM,
+};
+
+/**
+    Representation of a pointer in memory to a Lua value.
+*/
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct LuaValueId(usize);
+
+impl From<&LuaValue<'_>> for LuaValueId {
+    fn from(value: &LuaValue<'_>) -> Self {
+        Self(value.to_pointer() as usize)
+    }
+}
+
+impl From<&LuaTable<'_>> for LuaValueId {
+    fn from(table: &LuaTable) -> Self {
+        Self(table.to_pointer() as usize)
+    }
+}
+
+/**
+    Formats the given value, recursively formatting tables
+    up to the maximum depth specified in the config.
+
+    NOTE: We return a result here but it's really just to make handling
+    of the `write!` calls easier. Writing into a string should never fail.
+*/
+pub(crate) fn format_value_recursive(
+    value: &LuaValue,
+    config: &ValueFormatConfig,
+    visited: &mut HashSet<LuaValueId>,
+    depth: usize,
+) -> Result<String, fmt::Error> {
+    let mut buffer = String::new();
+
+    if let LuaValue::Table(ref t) = value {
+        if depth >= config.max_depth {
+            write!(buffer, "{}", STYLE_DIM.apply_to("{ ... }"))?;
+        } else if !visited.insert(LuaValueId::from(t)) {
+            write!(buffer, "{}", STYLE_DIM.apply_to("{ recursive }"))?;
+        } else {
+            writeln!(buffer, "{}", STYLE_DIM.apply_to("{"))?;
+
+            for res in t.clone().pairs::<LuaValue, LuaValue>() {
+                let (key, value) = res.expect("conversion to LuaValue should never fail");
+                let formatted = if let Some(plain_key) = lua_value_as_plain_string_key(&key) {
+                    format!(
+                        "{plain_key} {} {}{}",
+                        STYLE_DIM.apply_to("="),
+                        format_value_recursive(&value, config, visited, depth + 1)?,
+                        STYLE_DIM.apply_to(","),
+                    )
+                } else {
+                    format!(
+                        "{}{}{} {} {}{}",
+                        STYLE_DIM.apply_to("["),
+                        format_value_recursive(&key, config, visited, depth + 1)?,
+                        STYLE_DIM.apply_to("]"),
+                        STYLE_DIM.apply_to("="),
+                        format_value_recursive(&value, config, visited, depth + 1)?,
+                        STYLE_DIM.apply_to(","),
+                    )
+                };
+                buffer.push_str(&formatted);
+            }
+
+            visited.remove(&LuaValueId::from(t));
+            writeln!(buffer, "{}", STYLE_DIM.apply_to("}"))?;
+        }
+    } else {
+        write!(buffer, "{}", format_value_styled(value))?;
+    }
+
+    Ok(buffer)
+}