1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use anstyle::*;
use anyhow::Result;
use clap::Parser;
use sp1_build::{build_program, BuildArgs};
use sp1_core::utils::{setup_logger, setup_tracer};
use sp1_prover::SP1Stdin;
use sp1_sdk::ProverClient;
use std::time::Instant;
use std::{env, fs::File, io::Read, path::PathBuf, str::FromStr};

use crate::util::{elapsed, write_status};

#[derive(Debug, Clone)]
enum Input {
    FilePath(PathBuf),
    HexBytes(Vec<u8>),
}

fn is_valid_hex_string(s: &str) -> bool {
    if s.len() % 2 != 0 {
        return false;
    }
    // All hex digits with optional 0x prefix
    s.starts_with("0x") && s[2..].chars().all(|c| c.is_ascii_hexdigit())
        || s.chars().all(|c| c.is_ascii_hexdigit())
}

impl FromStr for Input {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if is_valid_hex_string(s) {
            // Remove 0x prefix if present
            let s = if s.starts_with("0x") {
                s.strip_prefix("0x").unwrap()
            } else {
                s
            };
            if s.is_empty() {
                return Ok(Input::HexBytes(Vec::new()));
            }
            if !s.chars().all(|c| c.is_ascii_hexdigit()) {
                return Err("Invalid hex string.".to_string());
            }
            let bytes = hex::decode(s).map_err(|e| e.to_string())?;
            Ok(Input::HexBytes(bytes))
        } else if PathBuf::from(s).exists() {
            Ok(Input::FilePath(PathBuf::from(s)))
        } else {
            Err("Input must be a valid file path or hex string.".to_string())
        }
    }
}

#[derive(Parser)]
#[command(name = "prove", about = "(default) Build and prove a program")]
pub struct ProveCmd {
    #[clap(long, value_parser)]
    input: Option<Input>,

    #[clap(long, action)]
    output: Option<PathBuf>,

    #[clap(long, action)]
    profile: bool,

    #[clap(long, action)]
    verbose: bool,

    #[clap(flatten)]
    build_args: BuildArgs,
}

impl ProveCmd {
    pub fn run(&self) -> Result<()> {
        let elf_path = build_program(&self.build_args, None)?;

        if !self.profile {
            match env::var("RUST_LOG") {
                Ok(_) => {}
                Err(_) => env::set_var("RUST_LOG", "info"),
            }
            setup_logger();
        } else {
            match env::var("RUST_TRACER") {
                Ok(_) => {}
                Err(_) => env::set_var("RUST_TRACER", "info"),
            }
            setup_tracer();
        }

        let mut elf = Vec::new();
        File::open(elf_path.as_path().as_str())
            .expect("failed to open input file")
            .read_to_end(&mut elf)
            .expect("failed to read from input file");

        let mut stdin = SP1Stdin::new();
        if let Some(ref input) = self.input {
            match input {
                Input::FilePath(ref path) => {
                    let mut file = File::open(path).expect("failed to open input file");
                    let mut bytes = Vec::new();
                    file.read_to_end(&mut bytes)?;
                    stdin.write_slice(&bytes);
                }
                Input::HexBytes(ref bytes) => {
                    stdin.write_slice(bytes);
                }
            }
        }

        let start_time = Instant::now();
        let client = ProverClient::new();
        let (pk, _) = client.setup(&elf);
        let proof = client.prove(&pk, stdin).run().unwrap();

        if let Some(ref path) = self.output {
            proof
                .save(path.to_str().unwrap())
                .expect("failed to save proof");
        }

        let elapsed = elapsed(start_time.elapsed());
        let green = AnsiColor::Green.on_default().effects(Effects::BOLD);
        write_status(
            &green,
            "Finished",
            format!("proving in {}", elapsed).as_str(),
        );

        Ok(())
    }
}