feat: Generic build and install phases, which only compile instead of proving

3 files changed, 44 insertions, 9 deletions

diff --git a/flake.nix b/flake.nix
index 7299b4b..9f315ac 100644
--- a/flake.nix
+++ b/flake.nix
@@ -20,7 +20,36 @@
     pkgs = import nixpkgs { system = system; overlays = [ mcl-blockchain.overlays.default ]; };
     callPackage = pkgs.lib.callPackageWith pkgs;
 
+    withZKVMPhases = currentPackage: with currentPackage; {
+        buildPhase = ''
+          root_dir="$PWD"
+          runHook preBuild
+          cd "$root_dir"
+
+          export INPUTS="$PWD/Vertices-010.in"
+          cargo build --bin ${hostBin} --release
+
+          runHook postBuild
+        '';
+
+        installPhase = ''
+          runHook preInstall
+          mkdir -p "$out"/bin
+          for bin in $(find . -type f -regex ".*release/[^/]*" -executable -print)
+          do
+            mv "$bin" "$out"/bin/
+          done
+          runHook postInstall
+        '';
+
+        doNotPostBuildInstallCargoBinaries = true;
+      } // currentPackage;
+
+    args-zkVM = {
+      craneLib-default = crane.mkLib pkgs;
+      inherit withZKVMPhases;
+    };
   in {
-    packages.${system}.risc0 = callPackage ./zkvms/risc0/default.nix {craneLib-default = crane.mkLib pkgs;};
+    packages.${system}.risc0 = callPackage ./zkvms/risc0/default.nix args-zkVM;
   };
 }
diff --git a/zkvms/risc0/default.nix b/zkvms/risc0/default.nix
index 42ebf8a..c491217 100644
--- a/zkvms/risc0/default.nix
+++ b/zkvms/risc0/default.nix
@@ -4,6 +4,7 @@
   metacraft-labs,
   pkg-config,
   craneLib-default,
+  withZKVMPhases,
 }:
 let
   fs = lib.fileset;
@@ -28,7 +29,7 @@ let
   craneLib = craneLib-default.overrideToolchain metacraft-labs.risc0;
   cargoArtifacts = craneLib.buildDepsOnly commonArgs;
 in
-  craneLib.buildPackage (commonArgs
+  craneLib.buildPackage (withZKVMPhases (commonArgs
     // {
       inherit cargoArtifacts;
 
@@ -44,11 +45,16 @@ in
       '';
 
       preBuild = ''
-        INPUTS="$PWD/Vertices-010.in"
-        export INPUTS
-        cd zkvms/risc0
-        just prove
+        cd zkvms/risc0/guest
+        cargo build --release --target riscv32im-risc0-zkvm-elf
+        ln -s ../../../../zkvms/risc0/guest/target/riscv32im-risc0-zkvm-elf/release/guest ../host/src/guest
+      '';
+
+      hostBin = "host-risc0";
+
+      postInstall = ''
+        ln -s "${metacraft-labs.risc0}"/bin/r0vm "$out"/bin/r0vm
       '';
 
       doCheck = false;
-    })
+    }))
diff --git a/zkvms/risc0/host/src/main.rs b/zkvms/risc0/host/src/main.rs
index de5417c..653f11d 100644
--- a/zkvms/risc0/host/src/main.rs
+++ b/zkvms/risc0/host/src/main.rs
@@ -1,6 +1,6 @@
 use risc0_zkvm::{default_prover, default_executor, ExecutorEnv};
 
-static HELLO_GUEST_ELF: &[u8] = include_bytes!("../../../../zkvms/risc0/guest/target/riscv32im-risc0-zkvm-elf/release/guest");
+static HELLO_GUEST_ELF: &[u8] = include_bytes!("./guest");
 
 fn main() {
     let args: Vec<String> = std::env::args().collect();
@@ -29,7 +29,7 @@ fn main() {
             let prover = default_prover();
 
             // Produce a receipt by proving the specified ELF binary.
-            let receipt = prover.prove(env, HELLO_GUEST_ELF).unwrap().receipt;
+            let receipt = prover.prove(env, HELLO_GUEST_ELF).expect("Error occured").receipt;
 
             // Extract journal of receipt
             let journal: bool = receipt.journal.decode().unwrap();