z39_solver/
solver.rs

1use std::path::PathBuf;
2use std::process::Stdio;
3use tokio::io::AsyncWriteExt;
4use tokio::process::Command;
5use tokio::time::{timeout, Duration};
6use anyhow::Result;
7
8const Z3_VERSION: &str = "4.16.0";
9
10/// Find or download the z3 binary.
11/// Search order:
12/// 1. Z3_BIN env var
13/// 2. Same directory as the z39 binary (distribution layout)
14/// 3. ~/.local/share/z39/z3 (auto-downloaded)
15/// 4. PATH lookup
16/// If not found anywhere, auto-download from GitHub releases (macOS/Windows)
17/// or build from source (Linux, since official binaries require glibc 2.38+).
18pub async fn find_or_download_z3() -> anyhow::Result<PathBuf> {
19    // 1. Z3_BIN env var
20    if let Ok(p) = std::env::var("Z3_BIN") {
21        let pb = PathBuf::from(&p);
22        if pb.exists() { return Ok(pb); }
23    }
24
25    // 2. Same directory as the z39 binary (distribution layout)
26    if let Ok(exe) = std::env::current_exe() {
27        if let Some(dir) = exe.parent() {
28            for name in ["z3", "z3.exe"] {
29                let candidate = dir.join(name);
30                if candidate.exists() { return Ok(candidate); }
31            }
32        }
33    }
34
35    // 3. ~/.local/share/z39/z3 (auto-downloaded)
36    let cache_dir = dirs::data_local_dir()
37        .unwrap_or_else(|| PathBuf::from("."))
38        .join("z39");
39    let cached = cache_dir.join("z3");
40    if cfg!(target_os = "windows") {
41        let cached_exe = cache_dir.join("z3.exe");
42        if cached_exe.exists() { return Ok(cached_exe); }
43    }
44    if cached.exists() { return Ok(cached); }
45
46    // 4. PATH lookup
47    for dir in std::env::var("PATH").unwrap_or_default().split(':') {
48        let p = PathBuf::from(dir).join("z3");
49        if p.exists() { return Ok(p); }
50    }
51
52    // Not found — auto-download
53    eprintln!("z39: z3 not found, auto-provisioning v{}...", Z3_VERSION);
54    provision_z3(&cache_dir).await?;
55    if cfg!(target_os = "windows") {
56        Ok(cache_dir.join("z3.exe"))
57    } else {
58        Ok(cache_dir.join("z3"))
59    }
60}
61
62/// Download Z3 binary or build from source depending on platform.
63async fn provision_z3(cache_dir: &PathBuf) -> anyhow::Result<()> {
64    std::fs::create_dir_all(cache_dir)?;
65
66    #[cfg(target_os = "macos")]
67    {
68        let arch = if cfg!(target_arch = "aarch64") { "arm64" } else { "x64" };
69        let filename = format!("z3-{Z3_VERSION}-{arch}-osx-15.7.3.zip");
70        let url = format!("https://github.com/Z3Prover/z3/releases/download/z3-{Z3_VERSION}/{filename}");
71        download_and_extract(&url, cache_dir).await
72    }
73
74    #[cfg(target_os = "windows")]
75    {
76        let filename = format!("z3-{Z3_VERSION}-x64-win.zip");
77        let url = format!("https://github.com/Z3Prover/z3/releases/download/z3-{Z3_VERSION}/{filename}");
78        download_and_extract(&url, cache_dir).await
79    }
80
81    #[cfg(target_os = "linux")]
82    {
83        // Official Z3 binaries require glibc 2.38+, many Linux distros ship older.
84        // Build from source for maximum compatibility.
85        build_z3_from_source(cache_dir).await
86    }
87
88    #[cfg(not(any(target_os = "macos", target_os = "windows", target_os = "linux")))]
89    {
90        anyhow::bail!("z39: unsupported platform for auto-provisioning; install z3 manually and set Z3_BIN")
91    }
92}
93
94async fn download_and_extract(url: &str, cache_dir: &PathBuf) -> anyhow::Result<()> {
95    eprintln!("z39: downloading {url}");
96    let response = reqwest::get(url).await?;
97    if !response.status().is_success() {
98        anyhow::bail!("download failed: {}", response.status());
99    }
100    let bytes = response.bytes().await?;
101    let tmpdir = tempfile::tempdir()?;
102
103    let zip_path = tmpdir.path().join("z3.zip");
104    std::fs::write(&zip_path, &bytes)?;
105
106    let archive = std::io::BufReader::new(std::fs::File::open(&zip_path)?);
107    let mut zip = zip::ZipArchive::new(archive)?;
108
109    let dest_name = if cfg!(target_os = "windows") { "z3.exe" } else { "z3" };
110    let dest = cache_dir.join(dest_name);
111
112    for i in 0..zip.len() {
113        let mut file = zip.by_index(i)?;
114        let name = file.name().to_string();
115        // Match z3 binary inside the archive (any path ending in /bin/z3 or /z3)
116        if name.ends_with("/bin/z3") || name.ends_with("/bin/z3.exe") || name == "z3" || name == "z3.exe" {
117            let mut buf = Vec::new();
118            std::io::Read::read_to_end(&mut file, &mut buf)?;
119            std::fs::write(&dest, &buf)?;
120            #[cfg(unix)]
121            {
122                use std::os::unix::fs::PermissionsExt;
123                std::fs::set_permissions(&dest, std::fs::Permissions::from_mode(0o755))?;
124            }
125            eprintln!("z39: z3 installed to {}", dest.display());
126            return Ok(());
127        }
128    }
129    anyhow::bail!("z3 binary not found in archive")
130}
131
132#[cfg(target_os = "linux")]
133async fn build_z3_from_source(cache_dir: &PathBuf) -> anyhow::Result<()> {
134    let src_url = format!("https://github.com/Z3Prover/z3/archive/refs/tags/z3-{Z3_VERSION}.tar.gz");
135    let tmpdir = tempfile::tempdir()?;
136
137    eprintln!("z39: downloading Z3 source...");
138    let response = reqwest::get(&src_url).await?;
139    if !response.status().is_success() {
140        anyhow::bail!("download failed: {}", response.status());
141    }
142    let bytes = response.bytes().await?;
143    let archive_path = tmpdir.path().join("z3.tar.gz");
144    std::fs::write(&archive_path, &bytes)?;
145
146    // Extract
147    let output = Command::new("tar")
148        .args(["-xzf", archive_path.to_str().unwrap(), "-C", tmpdir.path().to_str().unwrap()])
149        .output().await?;
150    if !output.status.success() {
151        anyhow::bail!("failed to extract Z3 source: {}", String::from_utf8_lossy(&output.stderr));
152    }
153
154    let src_dir = tmpdir.path().join(format!("z3-z3-{Z3_VERSION}"));
155    let build_dir = src_dir.join("build");
156    std::fs::create_dir_all(&build_dir)?;
157
158    // Configure
159    eprintln!("z39: configuring Z3...");
160    let output = Command::new("python3")
161        .args(["../scripts/mk_make.py"])
162        .current_dir(&build_dir)
163        .output().await?;
164    if !output.status.success() {
165        anyhow::bail!("Z3 configure failed: {}", String::from_utf8_lossy(&output.stderr));
166    }
167
168    // Build
169    let nproc = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
170    eprintln!("z39: building Z3 with {nproc} jobs (this may take a few minutes)...");
171    let output = Command::new("make")
172        .args(["-j", &nproc.to_string()])
173        .current_dir(&build_dir)
174        .output().await?;
175    if !output.status.success() {
176        anyhow::bail!("Z3 build failed: {}", String::from_utf8_lossy(&output.stderr));
177    }
178
179    // Copy to cache
180    let built = build_dir.join("z3");
181    let dest = cache_dir.join("z3");
182    std::fs::copy(&built, &dest)?;
183    use std::os::unix::fs::PermissionsExt;
184    std::fs::set_permissions(&dest, std::fs::Permissions::from_mode(0o755))?;
185
186    eprintln!("z39: Z3 v{} built and installed to {}", Z3_VERSION, dest.display());
187    Ok(())
188}
189
190// === Solver ===
191
192#[derive(Debug, Clone, PartialEq)]
193pub enum SolveStatus {
194    Sat,
195    Unsat,
196    Unknown,
197    Timeout,
198    Error(String),
199}
200
201#[derive(Debug, Clone)]
202pub struct SolveResult {
203    pub status: SolveStatus,
204    pub model: Option<String>,
205    pub raw_output: String,
206    pub duration_ms: u64,
207}
208
209impl SolveResult {
210    /// Token-optimized compact output
211    pub fn to_compact(&self) -> String {
212        let t = format!("{:.1}s", self.duration_ms as f64 / 1000.0);
213        match &self.status {
214            SolveStatus::Sat => {
215                let m = self.model.as_deref().unwrap_or("");
216                if m.is_empty() { format!("sat {t}") } else { format!("sat {m} {t}") }
217            }
218            SolveStatus::Unsat => format!("unsat {t}"),
219            SolveStatus::Unknown => format!("unknown {t}"),
220            SolveStatus::Timeout => format!("timeout {t}"),
221            SolveStatus::Error(e) => format!("error {e} {t}"),
222        }
223    }
224
225    pub fn is_sat(&self) -> bool { self.status == SolveStatus::Sat }
226    pub fn is_unsat(&self) -> bool { self.status == SolveStatus::Unsat }
227}
228
229/// Run z3 on SMT-LIB2 input
230pub async fn solve(z3_bin: &PathBuf, smt_input: &str, timeout_secs: u64) -> SolveResult {
231    let start = std::time::Instant::now();
232    let dur = if timeout_secs > 0 {
233        Duration::from_secs(timeout_secs)
234    } else {
235        Duration::from_secs(300)
236    };
237
238    let result = timeout(dur, run_z3_process(z3_bin, smt_input)).await;
239    let elapsed = start.elapsed().as_millis() as u64;
240
241    match result {
242        Ok(Ok(output)) => parse_z3_output(&output, elapsed),
243        Ok(Err(e)) => SolveResult {
244            status: SolveStatus::Error(e.to_string()),
245            model: None,
246            raw_output: String::new(),
247            duration_ms: elapsed,
248        },
249        Err(_) => SolveResult {
250            status: SolveStatus::Timeout,
251            model: None,
252            raw_output: String::new(),
253            duration_ms: elapsed,
254        },
255    }
256}
257
258async fn run_z3_process(z3_bin: &PathBuf, smt_input: &str) -> Result<String> {
259    let mut child = Command::new(z3_bin)
260        .arg("-in")
261        .stdin(Stdio::piped())
262        .stdout(Stdio::piped())
263        .stderr(Stdio::piped())
264        .spawn()?;
265
266    if let Some(mut stdin) = child.stdin.take() {
267        stdin.write_all(smt_input.as_bytes()).await?;
268        stdin.shutdown().await?;
269    }
270
271    let output = child.wait_with_output().await?;
272    Ok(String::from_utf8_lossy(&output.stdout).to_string())
273}
274
275fn parse_z3_output(raw: &str, elapsed_ms: u64) -> SolveResult {
276    let trimmed = raw.trim();
277    let first_line = trimmed.lines().next().unwrap_or("").trim();
278
279    let (status, model) = if first_line.starts_with("sat") || first_line == "sat" {
280        (SolveStatus::Sat, extract_model(trimmed))
281    } else if first_line.starts_with("unsat") {
282        (SolveStatus::Unsat, None)
283    } else if first_line.starts_with("unknown") {
284        (SolveStatus::Unknown, None)
285    } else if first_line.starts_with("error") {
286        (SolveStatus::Error(first_line.to_string()), None)
287    } else {
288        (SolveStatus::Error(format!("unexpected: {first_line}")), None)
289    };
290
291    SolveResult { status, model, raw_output: trimmed.to_string(), duration_ms: elapsed_ms }
292}
293
294fn extract_model(output: &str) -> Option<String> {
295    let mut assignments = Vec::new();
296    let lines: Vec<&str> = output.lines().collect();
297    for i in 0..lines.len() {
298        let line = lines[i].trim();
299        if line.starts_with("(define-fun") {
300            let parts: Vec<&str> = line.split_whitespace().collect();
301            if parts.len() >= 4 {
302                let name = parts[1];
303                if parts.len() >= 5 {
304                    let val = parts[parts.len() - 1].trim_end_matches(')');
305                    if !val.is_empty() && val != "(" && val != "()" {
306                        assignments.push(format!("{name}={val}"));
307                        continue;
308                    }
309                }
310                if i + 1 < lines.len() {
311                    let next = lines[i + 1].trim();
312                    let val = next.trim_end_matches(')');
313                    if !val.is_empty() && val != "(" {
314                        assignments.push(format!("{name}={val}"));
315                    }
316                }
317            }
318        }
319    }
320    if assignments.is_empty() { None } else { Some(assignments.join(" ")) }
321}
z39_solver/solver.rs

z39_solver/
solver.rs
