From 083b2d47c4192b24572d280b8367cf1eeb5caa62 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matthias=20G=C3=B6rgens?= <matthias.goergens@gmail.com>
Date: Thu, 9 Jan 2025 19:02:40 +0800
Subject: [PATCH] Clean up 'hints' for guests (#821)

Add the functionality to chain writes to the hints stream.
---
 ceno_host/src/lib.rs        |  5 +--
 ceno_host/tests/test_elf.rs | 61 +++++++++++++++----------------------
 2 files changed, 28 insertions(+), 38 deletions(-)

diff --git a/ceno_host/src/lib.rs b/ceno_host/src/lib.rs
index d6cd94ea1..fe6765c27 100644
--- a/ceno_host/src/lib.rs
+++ b/ceno_host/src/lib.rs
@@ -120,14 +120,15 @@ impl From<&CenoStdin> for Vec<u32> {
 }
 
 impl CenoStdin {
-    pub fn write_slice(&mut self, bytes: AlignedVec) {
+    pub fn write_slice(&mut self, bytes: AlignedVec) -> &mut Self {
         self.items.push(bytes);
+        self
     }
 
     pub fn write(
         &mut self,
         item: &impl for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, Error>>,
-    ) -> Result<(), Error> {
+    ) -> Result<&mut Self, Error> {
         to_bytes::<Error>(item).map(|bytes| self.write_slice(bytes))
     }
 }
diff --git a/ceno_host/tests/test_elf.rs b/ceno_host/tests/test_elf.rs
index 39d132751..97fcd751b 100644
--- a/ceno_host/tests/test_elf.rs
+++ b/ceno_host/tests/test_elf.rs
@@ -7,6 +7,7 @@ use ceno_emul::{
 };
 use ceno_host::CenoStdin;
 use itertools::{Itertools, enumerate, izip};
+use rand::{Rng, thread_rng};
 use tiny_keccak::keccakf;
 
 #[test]
@@ -111,14 +112,15 @@ fn test_ceno_rt_io() -> Result<()> {
 
 #[test]
 fn test_hints() -> Result<()> {
-    let mut hints = CenoStdin::default();
-    hints.write(&true)?;
-    hints.write(&"This is my hint string.".to_string())?;
-    hints.write(&1997_u32)?;
-    hints.write(&1999_u32)?;
-
-    let all_messages =
-        messages_to_strings(&ceno_host::run(CENO_PLATFORM, ceno_examples::hints, &hints));
+    let all_messages = messages_to_strings(&ceno_host::run(
+        CENO_PLATFORM,
+        ceno_examples::hints,
+        CenoStdin::default()
+            .write(&true)?
+            .write(&"This is my hint string.".to_string())?
+            .write(&1997_u32)?
+            .write(&1999_u32)?,
+    ));
     for (i, msg) in enumerate(&all_messages) {
         println!("{i}: {msg}");
     }
@@ -128,17 +130,12 @@ fn test_hints() -> Result<()> {
 
 #[test]
 fn test_bubble_sorting() -> Result<()> {
-    use rand::Rng;
-    let mut hints = CenoStdin::default();
-    let mut rng = rand::thread_rng();
-
-    // Provide some random numbers to sort.
-    hints.write(&(0..1_000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>())?;
-
+    let mut rng = thread_rng();
     let all_messages = messages_to_strings(&ceno_host::run(
         CENO_PLATFORM,
         ceno_examples::quadratic_sorting,
-        &hints,
+        // Provide some random numbers to sort.
+        CenoStdin::default().write(&(0..1_000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>())?,
     ));
     for msg in &all_messages {
         print!("{msg}");
@@ -147,17 +144,12 @@ fn test_bubble_sorting() -> Result<()> {
 }
 #[test]
 fn test_sorting() -> Result<()> {
-    use rand::Rng;
-    let mut hints = CenoStdin::default();
-    let mut rng = rand::thread_rng();
-
-    // Provide some random numbers to sort.
-    hints.write(&(0..1000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>())?;
-
+    let mut rng = thread_rng();
     let all_messages = messages_to_strings(&ceno_host::run(
         CENO_PLATFORM,
         ceno_examples::sorting,
-        &hints,
+        // Provide some random numbers to sort.
+        CenoStdin::default().write(&(0..1000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>())?,
     ));
     for (i, msg) in enumerate(&all_messages) {
         println!("{i}: {msg}");
@@ -167,9 +159,8 @@ fn test_sorting() -> Result<()> {
 
 #[test]
 fn test_median() -> Result<()> {
-    use rand::Rng;
     let mut hints = CenoStdin::default();
-    let mut rng = rand::thread_rng();
+    let mut rng = thread_rng();
 
     // Provide some random numbers to find the median of.
     let mut nums = (0..1000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>();
@@ -192,23 +183,22 @@ fn test_median() -> Result<()> {
 #[test]
 #[should_panic(expected = "Trap IllegalInstruction")]
 fn test_hashing_fail() {
-    use rand::Rng;
-    let mut hints = CenoStdin::default();
-    let mut rng = rand::thread_rng();
+    let mut rng = thread_rng();
 
     let mut nums = (0..1_000).map(|_| rng.gen::<u32>()).collect::<Vec<_>>();
     // Add a duplicate number to make uniqueness check fail:
     nums[211] = nums[907];
-    hints.write(&nums).unwrap();
 
-    let _ = ceno_host::run(CENO_PLATFORM, ceno_examples::hashing, &hints);
+    let _ = ceno_host::run(
+        CENO_PLATFORM,
+        ceno_examples::hashing,
+        CenoStdin::default().write(&nums).unwrap(),
+    );
 }
 
 #[test]
 fn test_hashing() -> Result<()> {
-    use rand::Rng;
-    let mut hints = CenoStdin::default();
-    let mut rng = rand::thread_rng();
+    let mut rng = thread_rng();
 
     // Provide some unique random numbers to verify:
     let uniques: Vec<u32> = {
@@ -219,11 +209,10 @@ fn test_hashing() -> Result<()> {
             .collect::<Vec<_>>()
     };
 
-    hints.write(&uniques)?;
     let all_messages = messages_to_strings(&ceno_host::run(
         CENO_PLATFORM,
         ceno_examples::hashing,
-        &hints,
+        CenoStdin::default().write(&uniques)?,
     ));
     assert!(!all_messages.is_empty());
     for (i, msg) in enumerate(&all_messages) {