Skip to main content

aprender_contracts_cli/commands/
extract.rs

1//! `pv extract-pytorch` — extract kernels from `PyTorch` source.
2
3use provable_contracts::extract;
4use std::path::Path;
5
6pub fn run(target: &str, output: Option<&Path>) -> Result<(), Box<dyn std::error::Error>> {
7    println!("pv extract-pytorch");
8    println!("==================\n");
9    println!("Source: {target}\n");
10
11    let kernel = extract::extract_from_pytorch(target)?;
12
13    println!("Function: {}", kernel.function_name);
14    println!("Arguments: {}", kernel.arguments.len());
15    println!("Equations: {}", kernel.equations.len());
16    for eq in &kernel.equations {
17        println!("  {}: {}", eq.name, eq.formula);
18        println!("    pre:  {:?}", eq.preconditions);
19        println!("    post: {:?}", eq.postconditions);
20    }
21
22    let yaml = extract::kernel_to_yaml(&kernel);
23
24    if let Some(out) = output {
25        std::fs::write(out, &yaml)?;
26        println!("\nContract written to: {}", out.display());
27    } else {
28        let default_path = format!(
29            "contracts/{}-v1.yaml",
30            kernel.function_name.replace('_', "-")
31        );
32        let p = Path::new(&default_path);
33        if let Some(parent) = p.parent() {
34            std::fs::create_dir_all(parent)?;
35        }
36        std::fs::write(&default_path, &yaml)?;
37        println!("\nContract written to: {default_path}");
38    }
39
40    Ok(())
41}