smt_scope/parsers/z3/
terms.rs

1use std::collections::hash_map::Entry;
2
3#[cfg(feature = "mem_dbg")]
4use mem_dbg::{MemDbg, MemSize};
5use typed_index_collections::TiSlice;
6
7use crate::{
8    error::Either,
9    items::{
10        InstProofLink, InstantiationKind, Meaning, ProofIdx, ProofStep, ProofStepKind, QuantIdx,
11        Quantifier, Term, TermId, TermIdToIdxMap, TermIdx, TermKind,
12    },
13    Error, FxHashMap, Result, StringTable, TiVec,
14};
15
16use super::bugs::TermsBug;
17
18pub trait HasTermId {
19    fn term_id(&self) -> TermId;
20}
21
22impl HasTermId for Term {
23    fn term_id(&self) -> TermId {
24        self.id
25    }
26}
27
28impl HasTermId for ProofStep {
29    fn term_id(&self) -> TermId {
30        self.id
31    }
32}
33
34#[cfg_attr(feature = "mem_dbg", derive(MemSize, MemDbg))]
35#[derive(Debug)]
36pub struct TermStorage<K: From<usize> + Copy, V: HasTermId> {
37    term_id_map: TermIdToIdxMap<K>,
38    terms: TiVec<K, V>,
39}
40
41impl<K: From<usize> + Copy, V: HasTermId> Default for TermStorage<K, V> {
42    fn default() -> Self {
43        Self {
44            term_id_map: TermIdToIdxMap::default(),
45            terms: TiVec::default(),
46        }
47    }
48}
49
50impl<K: From<usize> + Copy, V: HasTermId> TermStorage<K, V> {
51    pub(super) fn new_term(&mut self, term: V) -> Result<K> {
52        self.terms.raw.try_reserve(1)?;
53        let id = term.term_id();
54        let idx = self.terms.push_and_get_key(term);
55        self.term_id_map.register_term(id, idx)?;
56        Ok(idx)
57    }
58
59    fn get_term(&self, term_id: TermId) -> Either<K, TermId> {
60        self.term_id_map
61            .get_term(&term_id)
62            .map_or(Either::Right(term_id), Either::Left)
63    }
64    pub(super) fn parse_id(
65        &self,
66        strings: &mut StringTable,
67        id: &str,
68    ) -> Result<Either<K, TermId>> {
69        let term_id = TermId::parse(strings, id)?;
70        Ok(self.get_term(term_id))
71    }
72    pub(super) fn parse_existing_id(&self, strings: &mut StringTable, id: &str) -> Result<K> {
73        self.parse_id(strings, id)?
74            .into_result()
75            .map_err(Error::UnknownId)
76    }
77
78    pub(super) fn terms(&self) -> &TiSlice<K, V> {
79        &self.terms
80    }
81
82    /// Perform a top-down dfs walk of the AST rooted at `idx` calling `f` on
83    /// each node. If `f` returns `None` then the walk is terminated early.
84    /// Otherwise the walk is restricted to the children returned by `f`.
85    pub fn ast_walk<T>(
86        &self,
87        idx: K,
88        mut f: impl FnMut(K, &V) -> core::result::Result<&[K], T>,
89    ) -> Option<T>
90    where
91        usize: From<K>,
92    {
93        let mut todo = vec![idx];
94        while let Some(idx) = todo.pop() {
95            let next = &self.terms[idx];
96            match f(idx, next) {
97                Ok(to_walk) => todo.extend_from_slice(to_walk),
98                Err(t) => return Some(t),
99            }
100        }
101        None
102    }
103
104    /// Perform a bottom-up dfs walk of the AST rooted at `idx` calling `d` on
105    /// each node while walking down and `u` on each node while walking up.
106    pub fn ast_walk_cached<'a, D, I>(
107        &self,
108        idx: K,
109        state: I,
110        cache: &'a mut FxHashMap<K, D>,
111        mut d: impl for<'c> FnMut(K, &'c V, &I) -> (&'c [K], I),
112        mut u: impl FnMut(K, &V, &FxHashMap<K, D>, &I) -> D,
113    ) -> &'a D
114    where
115        usize: From<K>,
116        K: Eq + core::hash::Hash,
117    {
118        fn trim<T: Eq + core::hash::Hash, V>(
119            node: &mut core::slice::Iter<T>,
120            cache: &FxHashMap<T, V>,
121        ) {
122            let slice = node.as_slice();
123            let nk = slice.iter().position(|idx| !cache.contains_key(idx));
124            *node = slice[nk.unwrap_or(slice.len())..].iter()
125        }
126
127        let mut todo = vec![(state, core::slice::from_ref(&idx).iter())];
128        while let Some((mut state, mut node)) = todo.pop() {
129            trim(&mut node, cache);
130            while let Some(idx) = node.as_slice().first().copied() {
131                let (new_node, new_state) = d(idx, &self.terms[idx], &state);
132                todo.push((state, node));
133                state = new_state;
134                node = new_node.iter();
135                trim(&mut node, cache);
136            }
137
138            let Some((state, node)) = todo.last_mut() else {
139                break;
140            };
141            let idx = node.next().copied().unwrap();
142            let next = &self.terms[idx];
143            let data = u(idx, next, &*cache, &*state);
144            let old = cache.insert(idx, data);
145            assert!(old.is_none());
146        }
147        &cache[&idx]
148    }
149}
150
151#[cfg_attr(feature = "mem_dbg", derive(MemSize, MemDbg))]
152#[derive(Debug, Default)]
153pub struct Terms {
154    pub(super) app_terms: TermStorage<TermIdx, Term>,
155    pub(super) proof_terms: TermStorage<ProofIdx, ProofStep>,
156    pub(super) named_asserts: NamedAsserts,
157
158    meanings: FxHashMap<TermIdx, Meaning>,
159    pub(super) bug: TermsBug,
160}
161
162impl Terms {
163    pub(crate) fn meaning(&self, tidx: TermIdx) -> Option<&Meaning> {
164        self.meanings.get(&tidx)
165    }
166
167    pub(super) fn get_instantiation_body(&self, inst: InstantiationKind) -> Option<TermIdx> {
168        let proved_term = match inst {
169            InstantiationKind::Axiom { body } => return Some(body),
170            InstantiationKind::NonAxiom { proof, .. } => match proof {
171                InstProofLink::HasProof(proof_idx) => {
172                    let proof = &self[proof_idx];
173                    if matches!(proof.kind, ProofStepKind::PR_QUANT_INST) {
174                        proof.result
175                    } else {
176                        return Some(proof.result);
177                    }
178                }
179                InstProofLink::ProofsDisabled(term_idx) => term_idx?,
180            },
181        };
182        // The proved term is of the form `quant-inst(¬(quant) ∨ (inst))`.
183        let proved_term = &self[proved_term];
184        assert_eq!(proved_term.child_ids.len(), 2);
185        Some(proved_term.child_ids[1])
186    }
187
188    pub(super) fn quant(&self, quant: TermIdx) -> Result<QuantIdx> {
189        self[quant]
190            .quant_idx()
191            .ok_or(Error::UnknownQuantifierIdx(quant))
192    }
193
194    pub(super) fn new_meaning(&mut self, mut tidx: TermIdx, meaning: Meaning) -> Result<TermIdx> {
195        self.meanings.try_reserve(1)?;
196        match self.meanings.entry(tidx) {
197            Entry::Occupied(old) => {
198                if old.get() != &meaning {
199                    let term = self.app_terms.terms[tidx].clone();
200                    tidx = self.app_terms.new_term(term)?;
201                    self.meanings.insert(tidx, meaning);
202                }
203            }
204            Entry::Vacant(empty) => {
205                empty.insert(meaning);
206            }
207        };
208        Ok(tidx)
209    }
210
211    /// Perform a top-down walk of the AST rooted at `tidx` calling `f` on each
212    /// node of kind `App` and walking the children that are returned. `Quant`
213    /// and `Var` nodes are skipped.
214    pub(super) fn app_walk<T>(
215        &self,
216        tidx: TermIdx,
217        mut f: impl FnMut(TermIdx, &Term) -> core::result::Result<&[TermIdx], T>,
218    ) -> core::result::Result<(), T> {
219        self.app_terms
220            .ast_walk(tidx, |tidx, term| match term.kind() {
221                TermKind::Var(_) => Ok(&[]),
222                TermKind::App(_) => f(tidx, term),
223                TermKind::Quant(_) => Ok(&[]),
224            })
225            .map_or(Ok(()), Err)
226    }
227
228    /// Heuristic to get body of instantiated quantifier. See documentation of
229    /// [`InstProofLink::ProofsDisabled`].
230    pub(super) fn last_term_from_instance(&self, strings: &StringTable) -> Option<TermIdx> {
231        let last_non_eq = self
232            .app_terms
233            .terms
234            .iter_enumerated()
235            .rev()
236            .find(|(_, term)| term.app_name().is_none_or(|name| &strings[*name] != "="));
237        let last_term = last_non_eq.filter(|(_, term)| {
238            term.app_name().is_some_and(|name| &strings[*name] == "or")
239                && term.child_ids.len() == 2
240                && {
241                    let neg_quant = &self[term.child_ids[0]];
242                    neg_quant
243                        .app_name()
244                        .is_some_and(|name| &strings[*name] == "not")
245                        && neg_quant.child_ids.len() == 1
246                        && self[neg_quant.child_ids[0]].quant_idx().is_some()
247                }
248        });
249        debug_assert!(
250            last_term.is_some(),
251            "{:?}",
252            last_non_eq.map(|(_, t)| t.app_name().map(|n| &strings[*n]))
253        );
254        last_term.map(|(idx, _)| idx)
255    }
256
257    pub fn is_true_const(&self, tidx: TermIdx) -> bool {
258        let id = self[tidx].id;
259        id.namespace.is_none() && id.id.is_some_and(|id| id.get() == 1)
260    }
261    pub fn is_false_const(&self, tidx: TermIdx) -> bool {
262        let id = self[tidx].id;
263        id.namespace.is_none() && id.id.is_some_and(|id| id.get() == 2)
264    }
265    pub fn is_bool_const(&self, tidx: TermIdx) -> bool {
266        self.is_true_const(tidx) || self.is_false_const(tidx)
267    }
268
269    /// Used only to give access to the `app_terms` field in `bugs.rs`.
270    pub(super) fn get_app_term_bug(&self, term_id: TermId) -> Result<TermIdx> {
271        self.app_terms
272            .get_term(term_id)
273            .into_result()
274            .map_err(Error::UnknownId)
275    }
276
277    pub(super) fn new_proof(
278        &mut self,
279        quants: &mut TiVec<QuantIdx, Quantifier>,
280        pidx: ProofIdx,
281        strings: &StringTable,
282    ) -> Result<()> {
283        let proof = &self[pidx];
284        if !proof.kind.is_asserted() {
285            return Ok(());
286        }
287        let ridx = proof.result;
288
289        // Mark any quantifiers in the asserted term as blaming this proof step.
290        self.app_terms.ast_walk::<super::Never>(ridx, |_, term| {
291            if term.has_var().is_some() {
292                return Ok(&[]);
293            }
294            if let TermKind::Quant(qidx) = term.kind() {
295                // TODO: how to handle multiple blames here? Currently we use
296                // the latest one only.
297                // debug_assert!(quants[qidx].blame.is_none());
298                quants[qidx].blame = Some(pidx);
299            }
300            Ok(&term.child_ids)
301        });
302
303        let result = &self[ridx];
304        match result.child_ids.len() {
305            0 => {
306                if let Some(implication) = self.named_asserts.seen.remove(&ridx) {
307                    self.named_asserts.named.try_reserve(2)?;
308                    let old = self.named_asserts.named.insert(implication, Some(pidx));
309                    debug_assert!(old.is_none());
310                    let old = self.named_asserts.named.insert(pidx, None);
311                    debug_assert!(old.is_none());
312                }
313            }
314            2 if result
315                .app_name()
316                .is_some_and(|name| &strings[*name] == "=>") =>
317            {
318                let lidx = result.child_ids[0];
319                let lhs = &self[lidx];
320                if lhs.app_name().is_some() && lhs.child_ids.is_empty() {
321                    self.named_asserts.seen.try_reserve(1)?;
322                    self.named_asserts.seen.insert(lidx, pidx);
323                }
324            }
325            _ => (),
326        };
327        Ok(())
328    }
329}
330
331/// Named assertions are required to get an unsat core. They are given to z3 as
332/// follows: `(assert (! my_assertion :named my_name))`. This is turned into:
333/// ```smt2
334/// (declare-const my_name Bool)
335/// (assert (=> my_name my_assertion))
336/// (assert my_name)
337/// ```
338/// Unfortunately from the log we cannot differentiate between the user writing
339/// the above directly or using `:named`, so we treat both the same (unlike z3
340/// which differentiates the two).
341///
342/// When we see the first line:\
343/// `[mk-proof] #_ asserted (=> my_name #_)`\
344/// we store it under `seen`. Then when we see the second line:\
345/// `[mk-proof] #_ asserted my_name`\
346/// we look it up in `seen` and move it to `named`.
347#[cfg_attr(feature = "mem_dbg", derive(MemSize, MemDbg))]
348#[derive(Debug, Default)]
349pub struct NamedAsserts {
350    pub seen: FxHashMap<TermIdx, ProofIdx>,
351    /// Mapping from `assertion -> Some(name)` or `name -> None`. That is,
352    /// if a `ProofIdx` is in this map it's either a named assertion or the
353    /// boolean name variable itself. To differentiate the two, the named
354    /// assertion variable has an entry in the hashmap with `Some(name)`, while
355    /// the name variable has an entry with `None`.
356    pub named: FxHashMap<ProofIdx, Option<ProofIdx>>,
357}
358
359impl std::ops::Index<TermIdx> for Terms {
360    type Output = Term;
361    fn index(&self, idx: TermIdx) -> &Self::Output {
362        &self.app_terms.terms[idx]
363    }
364}
365
366impl std::ops::IndexMut<TermIdx> for Terms {
367    fn index_mut(&mut self, idx: TermIdx) -> &mut Self::Output {
368        &mut self.app_terms.terms[idx]
369    }
370}
371
372impl std::ops::Index<ProofIdx> for Terms {
373    type Output = ProofStep;
374    fn index(&self, idx: ProofIdx) -> &Self::Output {
375        &self.proof_terms.terms[idx]
376    }
377}
378
379impl std::ops::IndexMut<ProofIdx> for Terms {
380    fn index_mut(&mut self, idx: ProofIdx) -> &mut Self::Output {
381        &mut self.proof_terms.terms[idx]
382    }
383}