Skip to main content

liquid_fixpoint/
z3_subprocess.rs

1//! Subprocess for interfacing with Z3. We only added enough features to support
2//! [`suggustion_solver.rs`]
3
4use std::{
5    io::{self, BufRead, BufReader, Read as _, Write as _},
6    process::{Child, ChildStdin, ChildStdout, Command, Stdio},
7    sync::{Arc, Mutex},
8    thread::JoinHandle,
9};
10
11use crate::sexp::{Atom, Parser, Sexp};
12
13#[derive(Debug)]
14pub(crate) enum Z3Error {
15    Io(io::Error),
16    Timeout,
17}
18
19impl From<io::Error> for Z3Error {
20    fn from(err: io::Error) -> Self {
21        Self::Io(err)
22    }
23}
24
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub(crate) enum SatStatus {
27    Sat,
28    Unsat,
29    Unknown,
30}
31
32struct Response {
33    stdout: String,
34    stderr: String,
35}
36
37pub(crate) struct Z3Subprocess {
38    child: Child,
39    stdin: Option<ChildStdin>,
40    stdout: BufReader<ChildStdout>,
41    stderr: Arc<Mutex<Vec<u8>>>,
42    stderr_reader: Option<JoinHandle<()>>,
43    request_id: usize,
44}
45
46impl Z3Subprocess {
47    pub(crate) fn new() -> io::Result<Self> {
48        let mut child = Command::new("z3")
49            .args(["-smt2", "-in"])
50            .stdin(Stdio::piped())
51            .stdout(Stdio::piped())
52            .stderr(Stdio::piped())
53            .spawn()?;
54        let stdin = child.stdin.take().expect("Z3 stdin is piped");
55        let stdout = child.stdout.take().expect("Z3 stdout is piped");
56        let mut child_stderr = child.stderr.take().expect("Z3 stderr is piped");
57        let stderr = Arc::new(Mutex::new(Vec::new()));
58        let stderr_buf = Arc::clone(&stderr);
59        let stderr_reader = std::thread::spawn(move || {
60            let mut chunk = [0; 4096];
61            loop {
62                match child_stderr.read(&mut chunk) {
63                    Ok(0) | Err(_) => break,
64                    Ok(len) => stderr_buf.lock().unwrap().extend_from_slice(&chunk[..len]),
65                }
66            }
67        });
68        Ok(Self {
69            child,
70            stdin: Some(stdin),
71            stdout: BufReader::new(stdout),
72            stderr,
73            stderr_reader: Some(stderr_reader),
74            request_id: 0,
75        })
76    }
77
78    fn request(&mut self, commands: &str) -> io::Result<Response> {
79        let marker = format!("__flux_z3_end_{}__", self.request_id);
80        self.request_id += 1;
81        self.stderr.lock().unwrap().clear();
82        let stdin = self.stdin.as_mut().expect("Z3 stdin is open");
83        writeln!(stdin, "{commands}")?;
84        writeln!(stdin, "(echo \"{marker}\")")?;
85        stdin.flush()?;
86
87        // Use a marker to find the end of Z3's output.
88        let mut response = String::new();
89        loop {
90            let mut line = String::new();
91            match self.stdout.read_line(&mut line)? {
92                0 => {
93                    let status = self.child.try_wait()?.and_then(|status| status.code());
94                    let stderr = self.take_stderr();
95                    panic!(
96                        "Z3 closed stdout before the response marker (status {status:?}); stdout: {response:?}; stderr: {stderr:?}"
97                    );
98                }
99                _ if line.trim() == marker => {
100                    return Ok(Response { stdout: response, stderr: self.take_stderr() });
101                }
102                _ => response.push_str(&line),
103            }
104        }
105    }
106
107    fn execute(&mut self, commands: &str) -> io::Result<()> {
108        let Response { stdout, stderr } = self.request(commands)?;
109        assert!(
110            stdout.trim().is_empty(),
111            "unexpected Z3 setup response; commands: {commands:?}; stdout: {stdout:?}; stderr: {stderr:?}"
112        );
113        Ok(())
114    }
115
116    /// Check satisfiability (setup is a string that is sent to Z3 to be done
117    /// before the check)
118    pub(crate) fn check_sat(&mut self, setup: &str) -> Result<SatStatus, Z3Error> {
119        let Response { stdout, stderr } =
120            self.request(&format!("(push)\n{setup}\n(check-sat)\n(pop)"))?;
121        match stdout.trim() {
122            "sat" => Ok(SatStatus::Sat),
123            "unsat" => Ok(SatStatus::Unsat),
124            "unknown" => Ok(SatStatus::Unknown),
125            _ if is_timeout(&stdout) => Err(Z3Error::Timeout),
126            _ => panic!("unexpected Z3 status response; stdout: {stdout:?}; stderr: {stderr:?}"),
127        }
128    }
129
130    /// Apply a tactic (as a string) under a context (setup), returning the
131    /// parsed formulas in each resulting goal.
132    pub(crate) fn apply(&mut self, setup: &str, tactic: &str) -> Result<Vec<Vec<Sexp>>, Z3Error> {
133        let Response { stdout, stderr } =
134            self.request(&format!("(push)\n{setup}\n(apply {tactic})\n(pop)"))?;
135        if is_timeout(&stdout) {
136            return Err(Z3Error::Timeout);
137        }
138        Ok(parse_goals(&stdout, &stderr))
139    }
140
141    pub(crate) fn scoped<R, E>(
142        &mut self,
143        setup: &str,
144        f: impl FnOnce(&mut Self) -> Result<R, E>,
145    ) -> Result<R, E>
146    where
147        E: From<Z3Error> + From<io::Error>,
148    {
149        self.execute(&format!("(push)\n{setup}"))?;
150        let result = f(self);
151        let pop = self.execute("(pop)");
152        match result {
153            Err(err) => Err(err),
154            Ok(value) => pop.map(|_| value).map_err(E::from),
155        }
156    }
157
158    pub(crate) fn reset(&mut self, declarations: &str) -> io::Result<()> {
159        self.execute(&format!("(reset)\n{declarations}"))
160    }
161
162    fn take_stderr(&self) -> String {
163        let bytes = std::mem::take(&mut *self.stderr.lock().unwrap());
164        String::from_utf8_lossy(&bytes).into_owned()
165    }
166}
167
168impl Drop for Z3Subprocess {
169    fn drop(&mut self) {
170        self.stdin.take();
171        let _ = self.child.wait();
172        if let Some(reader) = self.stderr_reader.take() {
173            let _ = reader.join();
174        }
175    }
176}
177
178fn is_timeout(stdout: &str) -> bool {
179    let output = stdout.trim().to_ascii_lowercase();
180    // This is to the best of my understanding what Z3 outputs.
181    output.starts_with("(error ") && output.contains("canceled")
182}
183
184fn parse_goals(stdout: &str, stderr: &str) -> Vec<Vec<Sexp>> {
185    let parsed = Parser::new(stdout).parse_all().unwrap_or_else(|err| {
186        panic!("failed to parse Z3 goals response: {err:?}; stdout: {stdout:?}; stderr: {stderr:?}")
187    });
188    let [Sexp::List(items)] = parsed.as_slice() else {
189        panic!("expected one `(goals ...)` response; stdout: {stdout:?}; stderr: {stderr:?}");
190    };
191    if atom_symbol(items.first()) != Some("goals") {
192        panic!(
193            "Z3 goals response does not start with `goals`; stdout: {stdout:?}; stderr: {stderr:?}"
194        );
195    }
196    items[1..]
197        .iter()
198        .map(|goal| {
199            let Sexp::List(entries) = goal else {
200                panic!("Z3 goal response entry is not a list; entry: {goal:?}; stdout: {stdout:?}; stderr: {stderr:?}");
201            };
202            if atom_symbol(entries.first()) != Some("goal") {
203                panic!("expected `(goal ...)` in Z3 goals response; entry: {goal:?}; stdout: {stdout:?}; stderr: {stderr:?}");
204            }
205            let metadata = entries[1..]
206                .iter()
207                .position(|entry| {
208                    atom_symbol(Some(entry)).is_some_and(|name| name.starts_with(':'))
209                })
210                .unwrap_or(entries.len() - 1);
211            entries[1..1 + metadata].to_vec()
212        })
213        .collect()
214}
215
216fn atom_symbol(sexp: Option<&Sexp>) -> Option<&str> {
217    match sexp {
218        Some(Sexp::Atom(Atom::S(symbol))) => Some(symbol),
219        _ => None,
220    }
221}