1 files changed, 86 insertions, 2 deletions

diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 7a6bbcafb30f..f9ada2b627b5 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -1,11 +1,13 @@
 //! This module implements the backing representation of runtime
 //! values in the Nix language.
 use std::cmp::Ordering;
+use std::num::{NonZeroI32, NonZeroUsize};
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::rc::Rc;
 use std::{cell::Ref, fmt::Display};
 
+use lexical_core::format::CXX_LITERAL;
 use serde::{Deserialize, Serialize};
 
 #[cfg(feature = "arbitrary")]
@@ -32,6 +34,8 @@ pub use thunk::Thunk;
 
 use self::thunk::ThunkSet;
 
+use lazy_static::lazy_static;
+
 #[warn(variant_size_differences)]
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(untagged)]
@@ -72,6 +76,17 @@ pub enum Value {
     UnresolvedPath(PathBuf),
 }
 
+lazy_static! {
+    static ref WRITE_FLOAT_OPTIONS: lexical_core::WriteFloatOptions =
+        lexical_core::WriteFloatOptionsBuilder::new()
+            .trim_floats(true)
+            .round_mode(lexical_core::write_float_options::RoundMode::Round)
+            .positive_exponent_break(Some(NonZeroI32::new(5).unwrap()))
+            .max_significant_digits(Some(NonZeroUsize::new(6).unwrap()))
+            .build()
+            .unwrap();
+}
+
 // Helper macros to generate the to_*/as_* macros while accounting for
 // thunks.
 
@@ -514,9 +529,78 @@ impl TotalDisplay for Value {
             Value::Builtin(builtin) => builtin.fmt(f),
 
             // Nix prints floats with a maximum precision of 5 digits
-            // only.
+            // only. Except when it decides to use scientific notation
+            // (with a + after the `e`, and zero-padded to 0 digits)
             Value::Float(num) => {
-                write!(f, "{}", format!("{:.5}", num).trim_end_matches(['.', '0']))
+                let mut buf = [b'0'; lexical_core::BUFFER_SIZE];
+                let mut s = lexical_core::write_with_options::<f64, { CXX_LITERAL }>(
+                    num.clone(),
+                    &mut buf,
+                    &WRITE_FLOAT_OPTIONS,
+                );
+
+                // apply some postprocessing on the buffer. If scientific
+                // notation is used (we see an `e`), and the next character is
+                // a digit, add the missing `+` sign.)
+                let mut new_s = Vec::with_capacity(s.len());
+
+                if s.contains(&b'e') {
+                    for (i, c) in s.iter().enumerate() {
+                        // encountered `e`
+                        if c == &b'e' {
+                            // next character is a digit (so no negative exponent)
+                            if s.len() > i && s[i + 1].is_ascii_digit() {
+                                // copy everything from the start up to (including) the e
+                                new_s.extend_from_slice(&s[0..=i]);
+                                // add the missing '+'
+                                new_s.push(b'+');
+                                // check for the remaining characters.
+                                // If it's only one, we need to prepend a trailing zero
+                                if s.len() == i + 2 {
+                                    new_s.push(b'0');
+                                }
+                                new_s.extend_from_slice(&s[i + 1..]);
+                                break;
+                            }
+                        }
+                    }
+
+                    // if we modified the scientific notation, flip the reference
+                    if new_s.len() != 0 {
+                        s = &mut new_s
+                    }
+                }
+                // else, if this is not scientific notation, and there's a
+                // decimal point, make sure we really drop trailing zeroes.
+                // In some cases, lexical_core doesn't.
+                else if s.contains(&b'.') {
+                    for (i, c) in s.iter().enumerate() {
+                        // at `.``
+                        if c == &b'.' {
+                            // trim zeroes from the right side.
+                            let frac = String::from_utf8_lossy(&s[i + 1..]);
+                            let frac_no_trailing_zeroes = frac.trim_end_matches("0");
+
+                            if frac.len() != frac_no_trailing_zeroes.len() {
+                                // we managed to strip something, construct new_s
+                                if frac_no_trailing_zeroes.is_empty() {
+                                    // if frac_no_trailing_zeroes is empty, the fractional part was all zeroes, so we can drop the decimal point as well
+                                    new_s.extend_from_slice(&s[0..=i - 1]);
+                                } else {
+                                    // else, assemble the rest of the string
+                                    new_s.extend_from_slice(&s[0..=i]);
+                                    new_s.extend_from_slice(frac_no_trailing_zeroes.as_bytes());
+                                }
+
+                                // flip the reference
+                                s = &mut new_s;
+                                break;
+                            }
+                        }
+                    }
+                }
+
+                write!(f, "{}", format!("{}", String::from_utf8_lossy(&s)))
             }
 
             // internal types