From 45dd06a7b5afd72523f1d24ef5438a4fb6c5d396 Mon Sep 17 00:00:00 2001
From: Sagar Dhawan <sdhawan@a16z.com>
Date: Tue, 22 Oct 2024 11:09:23 -0500
Subject: [PATCH] fix: more verbose installs and clean up downloaded archive

---
 jolt-core/src/host/toolchain.rs | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/jolt-core/src/host/toolchain.rs b/jolt-core/src/host/toolchain.rs
index c2f9c2c50..e0f75c2a4 100644
--- a/jolt-core/src/host/toolchain.rs
+++ b/jolt-core/src/host/toolchain.rs
@@ -28,8 +28,13 @@ pub fn install_toolchain() -> Result<()> {
             download_toolchain(&client, &toolchain_url)
         }))?;
         unpack_toolchain()?;
+        remove_archive()?;
         link_toolchain()?;
         write_tag_file()?;
+        println!(
+            "\"riscv32im-jolt-zkvm-elf\" toolchain installed successfully at {:?}",
+            jolt_dir()
+        );
     }
     Ok(())
 }
@@ -39,7 +44,6 @@ pub fn install_no_std_toolchain() -> Result<()> {
     std::process::Command::new("rustup")
         .args(["target", "add", "riscv32im-unknown-none-elf"])
         .output()?;
-
     Ok(())
 }
 
@@ -114,7 +118,6 @@ async fn download_toolchain(client: &Client, url: &str) -> Result<()> {
         fs::create_dir(&jolt_dir)?;
     }
 
-    println!("Downloading toolchain");
     let mut response = client.get(url).send().await?;
     if response.status().is_success() {
         let mut file = File::create(output_path)?;
@@ -146,6 +149,14 @@ async fn download_toolchain(client: &Client, url: &str) -> Result<()> {
     }
 }
 
+fn remove_archive() -> Result<()> {
+    let toolchain_archive = jolt_dir().join("rust-toolchain.tar.gz");
+    if toolchain_archive.exists() {
+        fs::remove_file(&toolchain_archive)?;
+    }
+    Ok(())
+}
+
 fn toolchain_url() -> String {
     let target = target_lexicon::HOST;
     format!(
@@ -154,6 +165,16 @@ fn toolchain_url() -> String {
     )
 }
 
+#[cfg(not(target_arch = "wasm32"))]
+pub fn uninstall_no_std_toolchain() -> Result<()> {
+    std::process::Command::new("rustup")
+        .args(["target", "remove", "riscv32im-unknown-none-elf"])
+        .output()?;
+
+    println!("\"riscv32im-unknown-none-elf\" toolchain uninstalled successfully");
+    Ok(())
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 /// Uninstalls the toolchain if it is already installed
 pub fn uninstall_toolchain() -> Result<()> {
@@ -181,10 +202,7 @@ pub fn uninstall_toolchain() -> Result<()> {
     }
 
     // Remove the downloaded toolchain archive
-    let toolchain_archive = jolt_dir().join("rust-toolchain.tar.gz");
-    if toolchain_archive.exists() {
-        fs::remove_file(&toolchain_archive)?;
-    }
+    remove_archive()?;
 
     // Remove the toolchain tag file
     let tag_file = toolchain_tag_file();
@@ -192,7 +210,7 @@ pub fn uninstall_toolchain() -> Result<()> {
         fs::remove_file(&tag_file)?;
     }
 
-    println!("Toolchain uninstalled successfully");
+    println!("\"riscv32im-jolt-zkvm-elf\" toolchain uninstalled successfully");
     Ok(())
 }