feat(tvix/eval): nonrecursive nix_cmp_ordering(), fixes b/339 r/7167

This commit rewrites Value::nix_cmp_ordering() into an equivalent
nonrecursive form.  Except for calls to Thunk::force(), the new form
no longer uses generators, and is async only because of the fact
that it calls Thunk::force().

I originally believed that this commit would make evaluation faster.
In fact it is slightly slower.  I believe this is due to the added
vec![] allocation.  I am investigating.

Prev-Nixpkgs-Benchmark: {"attrpath":"pkgsCross.aarch64-multiplatform.hello.outPath","peak-kbytes":"460048","system-seconds":"0.68","user-seconds":"5.73"}
This-Nixpkgs-Benchmark: {"attrpath":"pkgsCross.aarch64-multiplatform.hello.outPath","peak-kbytes":"460224","system-seconds":"0.67","user-seconds":"5.84"}
Change-Id: Ic627bc220d9c5aa3c5e68b9b8bf199837cd55af5
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10212
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI
Autosubmit: Adam Joseph <adam@westernsemico.com>

7 files changed, 75 insertions, 44 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 35b7549ffe..1a34ea5dfb 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -610,8 +610,9 @@ mod pure_builtins {
 
     #[builtin("lessThan")]
     async fn builtin_less_than(co: GenCo, x: Value, y: Value) -> Result<Value, ErrorKind> {
+        let span = generators::request_span(&co).await;
         Ok(Value::Bool(matches!(
-            x.nix_cmp_ordering(y, co).await?,
+            x.nix_cmp_ordering(y, co, span).await?,
             Ordering::Less
         )))
     }
diff --git a/tvix/eval/src/tests/tvix_tests/notyetpassing/eval-okay-list-comparison.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-list-comparison.exp
index c508d5366f..c508d5366f 100644
--- a/tvix/eval/src/tests/tvix_tests/notyetpassing/eval-okay-list-comparison.exp
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-list-comparison.exp
diff --git a/tvix/eval/src/tests/tvix_tests/notyetpassing/eval-okay-list-comparison.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-list-comparison.nix
index 79cc27be65..79cc27be65 100644
--- a/tvix/eval/src/tests/tvix_tests/notyetpassing/eval-okay-list-comparison.nix
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-list-comparison.nix
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index c0bdc5fc5e..7b64e9fc8f 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -375,6 +375,16 @@ impl Value {
         }
     }
 
+    pub(crate) async fn nix_eq_owned_genco(
+        self,
+        other: Value,
+        co: GenCo,
+        ptr_eq: PointerEquality,
+        span: LightSpan,
+    ) -> Result<Value, ErrorKind> {
+        self.nix_eq(other, &co, ptr_eq, span).await
+    }
+
     /// Compare two Nix values for equality, forcing nested parts of the structure
     /// as needed.
     ///
@@ -388,7 +398,7 @@ impl Value {
     pub(crate) async fn nix_eq(
         self,
         other: Value,
-        co: GenCo,
+        co: &GenCo,
         ptr_eq: PointerEquality,
         span: LightSpan,
     ) -> Result<Value, ErrorKind> {
@@ -416,13 +426,13 @@ impl Value {
                         }
                     };
 
-                    Thunk::force_(thunk, &co, span.clone()).await?
+                    Thunk::force_(thunk, co, span.clone()).await?
                 }
 
                 _ => a,
             };
 
-            let b = b.force(&co, span.clone()).await?;
+            let b = b.force(co, span.clone()).await?;
 
             debug_assert!(!matches!(a, Value::Thunk(_)));
             debug_assert!(!matches!(b, Value::Thunk(_)));
@@ -471,8 +481,8 @@ impl Value {
                     #[allow(clippy::single_match)] // might need more match arms later
                     match (a1.select("type"), a2.select("type")) {
                         (Some(v1), Some(v2)) => {
-                            let s1 = v1.clone().force(&co, span.clone()).await?.to_str();
-                            let s2 = v2.clone().force(&co, span.clone()).await?.to_str();
+                            let s1 = v1.clone().force(co, span.clone()).await?.to_str();
+                            let s2 = v2.clone().force(co, span.clone()).await?.to_str();
 
                             if let (Ok(s1), Ok(s2)) = (s1, s2) {
                                 if s1.as_str() == "derivation" && s2.as_str() == "derivation" {
@@ -490,10 +500,10 @@ impl Value {
 
                                     let result = out1
                                         .clone()
-                                        .force(&co, span.clone())
+                                        .force(co, span.clone())
                                         .await?
                                         .to_str()?
-                                        == out2.clone().force(&co, span.clone()).await?.to_str()?;
+                                        == out2.clone().force(co, span.clone()).await?.to_str()?;
                                     if !result {
                                         return Ok(Value::Bool(false));
                                     } else {
@@ -597,55 +607,71 @@ impl Value {
     gen_is!(is_bool, Value::Bool(_));
     gen_is!(is_attrs, Value::Attrs(_));
 
-    // TODO(amjoseph): de-asyncify this (when called directly by the VM)
     /// Compare `self` against other using (fallible) Nix ordering semantics.
     ///
     /// The function is intended to be used from within other generator
     /// functions or `gen!` blocks.
-    pub async fn nix_cmp_ordering(self, other: Self, co: GenCo) -> Result<Ordering, ErrorKind> {
-        Self::nix_cmp_ordering_(self, other, co).await
+    pub async fn nix_cmp_ordering(
+        self,
+        other: Self,
+        co: GenCo,
+        span: LightSpan,
+    ) -> Result<Ordering, ErrorKind> {
+        Self::nix_cmp_ordering_(self, other, co, span).await
     }
 
     async fn nix_cmp_ordering_(
-        mut myself: Self,
-        mut other: Self,
+        myself: Self,
+        other: Self,
         co: GenCo,
+        span: LightSpan,
     ) -> Result<Ordering, ErrorKind> {
-        'outer: loop {
-            match (myself, other) {
+        // this is a stack of ((v1,v2),peq) triples to be compared;
+        // after each triple is popped off of the stack, v1 is
+        // compared to v2 using peq-mode PointerEquality
+        let mut vals = vec![((myself, other), PointerEquality::ForbidAll)];
+
+        loop {
+            let ((mut a, mut b), ptr_eq) = if let Some(abp) = vals.pop() {
+                abp
+            } else {
+                // stack is empty, so they are equal
+                return Ok(Ordering::Equal);
+            };
+            if ptr_eq == PointerEquality::AllowAll {
+                if a.clone()
+                    .nix_eq(b.clone(), &co, PointerEquality::AllowAll, span.clone())
+                    .await?
+                    .as_bool()?
+                {
+                    continue;
+                }
+                a = a.force(&co, span.clone()).await?;
+                b = b.force(&co, span.clone()).await?;
+            }
+            let result = match (a, b) {
                 // same types
-                // Nix does not support NaN or "negative infinity",
-                // so its floats are in fact totally ordered.
-                (Value::Integer(i1), Value::Integer(i2)) => return Ok(i1.cmp(&i2)),
-                (Value::Float(f1), Value::Float(f2)) => return Ok(f1.total_cmp(&f2)),
-                (Value::String(s1), Value::String(s2)) => return Ok(s1.cmp(&s2)),
+                (Value::Integer(i1), Value::Integer(i2)) => i1.cmp(&i2),
+                (Value::Float(f1), Value::Float(f2)) => f1.total_cmp(&f2),
+                (Value::String(s1), Value::String(s2)) => s1.cmp(&s2),
                 (Value::List(l1), Value::List(l2)) => {
-                    for i in 0.. {
-                        if i == l2.len() {
-                            return Ok(Ordering::Greater);
-                        } else if i == l1.len() {
-                            return Ok(Ordering::Less);
-                        } else if !generators::check_equality(
-                            &co,
-                            l1[i].clone(),
-                            l2[i].clone(),
-                            PointerEquality::AllowAll,
-                        )
-                        .await?
-                        {
-                            // TODO: do we need to control `top_level` here?
-                            myself = generators::request_force(&co, l1[i].clone()).await;
-                            other = generators::request_force(&co, l2[i].clone()).await;
-                            continue 'outer;
+                    let max = l1.len().max(l2.len());
+                    for j in 0..max {
+                        let i = max - 1 - j;
+                        if i >= l2.len() {
+                            vals.push(((1.into(), 0.into()), PointerEquality::ForbidAll));
+                        } else if i >= l1.len() {
+                            vals.push(((0.into(), 1.into()), PointerEquality::ForbidAll));
+                        } else {
+                            vals.push(((l1[i].clone(), l2[i].clone()), PointerEquality::AllowAll));
                         }
                     }
-
-                    unreachable!()
+                    continue;
                 }
 
                 // different types
-                (Value::Integer(i1), Value::Float(f2)) => return Ok((i1 as f64).total_cmp(&f2)),
-                (Value::Float(f1), Value::Integer(i2)) => return Ok(f1.total_cmp(&(i2 as f64))),
+                (Value::Integer(i1), Value::Float(f2)) => (i1 as f64).total_cmp(&f2),
+                (Value::Float(f1), Value::Integer(i2)) => f1.total_cmp(&(i2 as f64)),
 
                 // unsupported types
                 (lhs, rhs) => {
@@ -654,6 +680,9 @@ impl Value {
                         rhs: rhs.type_of(),
                     })
                 }
+            };
+            if result != Ordering::Equal {
+                return Ok(result);
             }
         }
     }
diff --git a/tvix/eval/src/vm/generators.rs b/tvix/eval/src/vm/generators.rs
index 3a256ec55d..4fff498fe7 100644
--- a/tvix/eval/src/vm/generators.rs
+++ b/tvix/eval/src/vm/generators.rs
@@ -337,7 +337,7 @@ impl<'o> VM<'o> {
                             let values = *values;
                             self.reenqueue_generator(name, span.clone(), generator);
                             self.enqueue_generator("nix_eq", span.clone(), |co| {
-                                values.0.nix_eq(values.1, co, ptr_eq, span)
+                                values.0.nix_eq_owned_genco(values.1, co, ptr_eq, span)
                             });
                             return Ok(false);
                         }
diff --git a/tvix/eval/src/vm/macros.rs b/tvix/eval/src/vm/macros.rs
index 8a536ee466..34e94bb5fa 100644
--- a/tvix/eval/src/vm/macros.rs
+++ b/tvix/eval/src/vm/macros.rs
@@ -42,7 +42,8 @@ macro_rules! cmp_op {
         async fn compare(a: Value, b: Value, co: GenCo) -> Result<Value, ErrorKind> {
             let a = generators::request_force(&co, a).await;
             let b = generators::request_force(&co, b).await;
-            let ordering = a.nix_cmp_ordering(b, co).await?;
+            let span = generators::request_span(&co).await;
+            let ordering = a.nix_cmp_ordering(b, co, span).await?;
             Ok(Value::Bool(cmp_op!(@order $op ordering)))
         }
 
diff --git a/tvix/eval/src/vm/mod.rs b/tvix/eval/src/vm/mod.rs
index 615d77c0e9..99a913c46a 100644
--- a/tvix/eval/src/vm/mod.rs
+++ b/tvix/eval/src/vm/mod.rs
@@ -616,7 +616,7 @@ impl<'o> VM<'o> {
                     let gen_span = frame.current_light_span();
                     self.push_call_frame(span, frame);
                     self.enqueue_generator("nix_eq", gen_span.clone(), |co| {
-                        a.nix_eq(b, co, PointerEquality::ForbidAll, gen_span)
+                        a.nix_eq_owned_genco(b, co, PointerEquality::ForbidAll, gen_span)
                     });
                     return Ok(false);
                 }