Refactor platform impls a bit (#261)

crates/dockerized/src/compiler.rs

@@ -1,8 +1,10 @@
 use crate::{
     CompilerKind,
-    image::{base_image, base_zkvm_image, compiler_zkvm_image},
+    image::{base_image, base_zkvm_image, compiler_zkvm_image, image_registry},
     util::{
-        docker::{DockerBuildCmd, DockerRunCmd, docker_image_exists, force_rebuild},
+        docker::{
+            DockerBuildCmd, DockerRunCmd, docker_image_exists, docker_pull_image, force_rebuild,
+        },
         workspace_dir,
     },
     zkVMKind,
@@ -34,9 +36,19 @@ fn build_compiler_image(zkvm_kind: zkVMKind) -> Result<(), Error> {
     let base_zkvm_image = base_zkvm_image(zkvm_kind, false);
     let compiler_zkvm_image = compiler_zkvm_image(zkvm_kind);
 
-    if !force_rebuild && docker_image_exists(&compiler_zkvm_image)? {
-        info!("Image {compiler_zkvm_image} exists, skip building");
-        return Ok(());
+    if !force_rebuild {
+        if docker_image_exists(&compiler_zkvm_image)? {
+            info!("Image {compiler_zkvm_image} exists, skip building");
+            return Ok(());
+        }
+
+        if image_registry().is_some()
+            && docker_pull_image(&compiler_zkvm_image).is_ok()
+            && docker_image_exists(&compiler_zkvm_image)?
+        {
+            info!("Image {compiler_zkvm_image} pulled, skip building");
+            return Ok(());
+        }
     }
 
     let workspace_dir = workspace_dir()?;

crates/dockerized/src/util/docker.rs

@@ -281,6 +281,25 @@ pub fn docker_container_exists(container_name: impl AsRef<str>) -> Result<bool,
     Ok(stdout.trim() == container_name.as_ref())
 }
 
+pub fn docker_pull_image(image: impl AsRef<str>) -> Result<(), CommonError> {
+    let mut cmd = Command::new("docker");
+    let output = cmd
+        .args(["image", "pull", image.as_ref()])
+        .stdout(Stdio::inherit())
+        .output()
+        .map_err(|err| CommonError::command(&cmd, err))?;
+
+    if !output.status.success() {
+        Err(CommonError::command_exit_non_zero(
+            &cmd,
+            output.status,
+            Some(&output),
+        ))?
+    }
+
+    Ok(())
+}
+
 pub fn docker_image_exists(image: impl AsRef<str>) -> Result<bool, CommonError> {
     let mut cmd = Command::new("docker");
     let output = cmd

crates/dockerized/src/zkvm.rs

@@ -1,11 +1,11 @@
 use crate::{
     compiler::SerializedProgram,
-    image::{base_image, base_zkvm_image, server_zkvm_image},
+    image::{base_image, base_zkvm_image, image_registry, server_zkvm_image},
     util::{
         cuda::cuda_arch,
         docker::{
             DockerBuildCmd, DockerRunCmd, docker_container_exists, docker_image_exists,
-            force_rebuild, stop_docker_container,
+            docker_pull_image, force_rebuild, stop_docker_container,
         },
         home_dir, workspace_dir,
     },
@@ -45,9 +45,19 @@ fn build_server_image(zkvm_kind: zkVMKind, gpu: bool) -> Result<(), Error> {
     let base_zkvm_image = base_zkvm_image(zkvm_kind, gpu);
     let server_zkvm_image = server_zkvm_image(zkvm_kind, gpu);
 
-    if !force_rebuild && docker_image_exists(&server_zkvm_image)? {
-        info!("Image {server_zkvm_image} exists, skip building");
-        return Ok(());
+    if !force_rebuild {
+        if docker_image_exists(&server_zkvm_image)? {
+            info!("Image {server_zkvm_image} exists, skip building");
+            return Ok(());
+        }
+
+        if image_registry().is_some()
+            && docker_pull_image(&server_zkvm_image).is_ok()
+            && docker_image_exists(&server_zkvm_image)?
+        {
+            info!("Image {server_zkvm_image} pulled, skip building");
+            return Ok(());
+        }
     }
 
     let workspace_dir = workspace_dir()?;

crates/zkvm/airbender/platform/src/lib.rs

@@ -30,7 +30,9 @@ impl Platform for AirbenderPlatform {
             "Maximum output size is 32 bytes, got {} bytes",
             output.len()
         );
-        let words = array::from_fn(|i| u32::from_le_bytes(array::from_fn(|j| output[4 * i + j])));
+        let words = array::from_fn(|i| {
+            u32::from_le_bytes(array::from_fn(|j| *output.get(4 * i + j).unwrap_or(&0)))
+        });
         riscv_common::zksync_os_finish_success(&words);
     }
 

crates/zkvm/pico/platform/Cargo.toml

@@ -7,13 +7,14 @@ license.workspace = true
 
 [dependencies]
 # Pico dependencies
-pico-sdk.workspace = true
+pico-sdk = { workspace = true, features = ["kb"] }
 
 # Local dependencies
 ere-platform-trait.workspace = true
 
 [features]
 default = []
+libm = ["pico-sdk/libm"]
 
 [lints]
 workspace = true

crates/zkvm/sp1/platform/Cargo.toml

@@ -7,7 +7,7 @@ license.workspace = true
 
 [dependencies]
 # SP1 dependencies
-sp1-zkvm.workspace = true
+sp1-zkvm = { workspace = true, features = ["lib"] }
 
 # Local dependencies
 ere-platform-trait.workspace = true

crates/zkvm/ziren/platform/Cargo.toml

@@ -13,7 +13,11 @@ zkm-zkvm.workspace = true
 ere-platform-trait.workspace = true
 
 [features]
-default = []
+default = ["libm", "bump"]
+bump = ["zkm-zkvm/bump"]
+embedded = ["zkm-zkvm/embedded"]
+libm = ["zkm-zkvm/libm"]
+verify = ["zkm-zkvm/verify"]
 
 [lints]
 workspace = true