feat(tazjin/generator-example): add an example for genawaiter crate r/5855

This is an experiment for tvix-eval.

Change-Id: Ic752b5b125cefefeb1343e38a70beb364478e6eb
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8131
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>

5 files changed, 260 insertions, 0 deletions

diff --git a/users/tazjin/generator-example/.gitignore b/users/tazjin/generator-example/.gitignore
new file mode 100644
index 000000000000..ea8c4bf7f35f
--- /dev/null
+++ b/users/tazjin/generator-example/.gitignore
@@ -0,0 +1 @@
+/target
diff --git a/users/tazjin/generator-example/Cargo.lock b/users/tazjin/generator-example/Cargo.lock
new file mode 100644
index 000000000000..a6f25ee39429
--- /dev/null
+++ b/users/tazjin/generator-example/Cargo.lock
@@ -0,0 +1,124 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "genawaiter"
+version = "0.99.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c86bd0361bcbde39b13475e6e36cb24c329964aa2611be285289d1e4b751c1a0"
+dependencies = [
+ "genawaiter-macro",
+ "genawaiter-proc-macro",
+ "proc-macro-hack",
+]
+
+[[package]]
+name = "genawaiter-macro"
+version = "0.99.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b32dfe1fdfc0bbde1f22a5da25355514b5e450c33a6af6770884c8750aedfbc"
+
+[[package]]
+name = "genawaiter-proc-macro"
+version = "0.99.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "784f84eebc366e15251c4a8c3acee82a6a6f427949776ecb88377362a9621738"
+dependencies = [
+ "proc-macro-error",
+ "proc-macro-hack",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "generator-example"
+version = "0.1.0"
+dependencies = [
+ "genawaiter",
+]
+
+[[package]]
+name = "proc-macro-error"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "18f33027081eba0a6d8aba6d1b1c3a3be58cbb12106341c2d5759fcd9b5277e7"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a5b4b77fdb63c1eca72173d68d24501c54ab1269409f6b672c85deb18af69de"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+ "syn-mid",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-hack"
+version = "0.5.20+deprecated"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc375e1527247fe1a97d8b7156678dfe7c1af2fc075c9a4db3690ecd2a148068"
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.51"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d727cae5b39d21da60fa540906919ad737832fe0b1c165da3a34d6548c849d6"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8856d8364d252a14d474036ea1358d63c9e6965c8e5c1885c18f73d70bff9c7b"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "syn"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f4064b5b16e03ae50984a5a8ed5d4f8803e6bc1fd170a3cda91a1be4b18e3f5"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn-mid"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "baa8e7560a164edb1621a55d18a0c59abf49d360f47aa7b821061dd7eea7fac9"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "84a22b9f218b40614adcb3f4ff08b703773ad44fa9423e4e0d346d5db86e4ebc"
+
+[[package]]
+name = "version_check"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
diff --git a/users/tazjin/generator-example/Cargo.toml b/users/tazjin/generator-example/Cargo.toml
new file mode 100644
index 000000000000..faf313973f5e
--- /dev/null
+++ b/users/tazjin/generator-example/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "generator-example"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+genawaiter = "0.99.1"
diff --git a/users/tazjin/generator-example/README.md b/users/tazjin/generator-example/README.md
new file mode 100644
index 000000000000..0bec13ee9a6b
--- /dev/null
+++ b/users/tazjin/generator-example/README.md
@@ -0,0 +1,11 @@
+generator-example
+=================
+
+This is an experiment with the [`genawaiter`][] crate, to see if it
+could be suitable for dealing with the execution flattening problem in
+Tvix.
+
+It constructs a dummy example that is similar to some of the problems
+we have in Tvix that require generator-like thunk forcing.
+
+[`genawaiter`]: https://docs.rs/genawaiter/latest/genawaiter/index.html
diff --git a/users/tazjin/generator-example/src/main.rs b/users/tazjin/generator-example/src/main.rs
new file mode 100644
index 000000000000..4aa931caf853
--- /dev/null
+++ b/users/tazjin/generator-example/src/main.rs
@@ -0,0 +1,115 @@
+use genawaiter::rc::{Co, Gen};
+use std::cell::RefCell;
+use std::future::Future;
+use std::pin::Pin;
+use std::rc::Rc;
+
+#[derive(Debug)]
+enum ValueRepr {
+    Int(i64),
+    Thunk((i64, i64)),
+}
+
+#[derive(Clone, Debug)]
+struct Value(Rc<RefCell<ValueRepr>>);
+
+impl Value {
+    fn force(&self) {
+        let mut inner = self.0.borrow_mut();
+        match *inner {
+            ValueRepr::Int(_) => return,
+            ValueRepr::Thunk((a, b)) => {
+                *inner = ValueRepr::Int(a + b);
+            }
+        }
+    }
+
+    fn is_forced(&self) -> bool {
+        matches!(*self.0.borrow(), ValueRepr::Int(_))
+    }
+
+    fn int(&self) -> i64 {
+        match *self.0.borrow() {
+            ValueRepr::Int(i) => i,
+            ValueRepr::Thunk(_) => panic!("unforced thunk!"),
+        }
+    }
+}
+
+impl From<i64> for Value {
+    fn from(value: i64) -> Self {
+        Value(Rc::new(RefCell::new(ValueRepr::Int(value))))
+    }
+}
+
+impl From<(i64, i64)> for Value {
+    fn from(value: (i64, i64)) -> Self {
+        Value(Rc::new(RefCell::new(ValueRepr::Thunk(value))))
+    }
+}
+
+async fn list_maker(values: Vec<Value>, co: Co<Value>) -> Vec<i64> {
+    let mut output: Vec<i64> = vec![];
+
+    for value in values {
+        if !value.is_forced() {
+            co.yield_(value.clone()).await;
+        }
+
+        output.push(value.int());
+    }
+
+    output
+}
+
+async fn list_reverser(values: Vec<Value>, co: Co<Value>) -> Vec<i64> {
+    let mut output = list_maker(values, co).await;
+    output.reverse();
+    output
+}
+
+struct Frame {
+    gen: Gen<Value, (), Pin<Box<dyn Future<Output = Vec<i64>>>>>,
+}
+
+fn pin_future(
+    f: impl Future<Output = Vec<i64>> + 'static,
+) -> Pin<Box<dyn Future<Output = Vec<i64>>>> {
+    Box::pin(f)
+}
+
+fn main() {
+    let mut frames: Vec<Frame> = vec![];
+
+    let values: Vec<Value> = vec![
+        42.into(),
+        (12, 54).into(),
+        4.into(),
+        (40, 2).into(),
+        2.into(),
+    ];
+    let second = values.clone();
+
+    frames.push(Frame {
+        gen: Gen::new(|co| pin_future(list_maker(values, co))),
+    });
+
+    frames.push(Frame {
+        gen: Gen::new(|co| pin_future(list_reverser(second, co))),
+    });
+
+    for (idx, mut frame) in frames.into_iter().enumerate() {
+        loop {
+            match frame.gen.resume() {
+                genawaiter::GeneratorState::Yielded(val) => {
+                    println!("yielded {:?} in frame {}", val, idx);
+                    val.force();
+                }
+                genawaiter::GeneratorState::Complete(list) => {
+                    println!("result {}: {:?}", idx, list);
+                    break;
+                }
+            }
+        }
+    }
+}