1use num_rational::Rational64;
14
15use oxiz_core::ast::TermId;
16use oxiz_core::sort::SortId;
17
18use crate::optimization::{OptimizationResult, Optimizer};
19use crate::z3_compat::{BV, Bool, Int, Real, SatResult, Z3Context};
20
21macro_rules! build {
24 ($ctx:expr, $method:ident $(, $arg:expr)* ) => {
25 $ctx.tm.borrow_mut().$method($($arg),*)
26 };
27}
28
29impl Real {
32 #[must_use]
34 pub fn gt(ctx: &Z3Context, lhs: &Real, rhs: &Real) -> Bool {
35 let id = build!(ctx, mk_gt, lhs.id, rhs.id);
36 Bool { id }
37 }
38
39 #[must_use]
41 pub fn ge(ctx: &Z3Context, lhs: &Real, rhs: &Real) -> Bool {
42 let id = build!(ctx, mk_ge, lhs.id, rhs.id);
43 Bool { id }
44 }
45
46 #[must_use]
48 pub fn neg(ctx: &Z3Context, arg: &Real) -> Real {
49 let id = build!(ctx, mk_neg, arg.id);
50 Real { id }
51 }
52
53 #[must_use]
55 pub fn div(ctx: &Z3Context, lhs: &Real, rhs: &Real) -> Real {
56 let id = build!(ctx, mk_div, lhs.id, rhs.id);
57 Real { id }
58 }
59
60 #[must_use]
62 pub fn from_i64(ctx: &Z3Context, value: i64) -> Real {
63 let id = build!(ctx, mk_real, Rational64::new(value, 1));
64 Real { id }
65 }
66}
67
68#[must_use]
74pub fn ite_bool(ctx: &Z3Context, cond: &Bool, then_branch: &Bool, else_branch: &Bool) -> Bool {
75 let id = build!(ctx, mk_ite, cond.id, then_branch.id, else_branch.id);
76 Bool { id }
77}
78
79#[must_use]
81pub fn ite_int(ctx: &Z3Context, cond: &Bool, then_branch: &Int, else_branch: &Int) -> Int {
82 let id = build!(ctx, mk_ite, cond.id, then_branch.id, else_branch.id);
83 Int { id }
84}
85
86#[must_use]
88pub fn ite_real(ctx: &Z3Context, cond: &Bool, then_branch: &Real, else_branch: &Real) -> Real {
89 let id = build!(ctx, mk_ite, cond.id, then_branch.id, else_branch.id);
90 Real { id }
91}
92
93#[must_use]
97pub fn ite_bv(ctx: &Z3Context, cond: &Bool, then_branch: &BV, else_branch: &BV) -> BV {
98 let width = then_branch.width;
99 let id = build!(ctx, mk_ite, cond.id, then_branch.id, else_branch.id);
100 BV { id, width }
101}
102
103#[must_use]
107pub fn distinct_int(ctx: &Z3Context, args: &[Int]) -> Bool {
108 let ids = args.iter().map(|x| x.id);
109 let id = build!(ctx, mk_distinct, ids);
110 Bool { id }
111}
112
113#[must_use]
115pub fn distinct_real(ctx: &Z3Context, args: &[Real]) -> Bool {
116 let ids = args.iter().map(|x| x.id);
117 let id = build!(ctx, mk_distinct, ids);
118 Bool { id }
119}
120
121#[must_use]
123pub fn distinct_bv(ctx: &Z3Context, args: &[BV]) -> Bool {
124 let ids = args.iter().map(|x| x.id);
125 let id = build!(ctx, mk_distinct, ids);
126 Bool { id }
127}
128
129#[derive(Debug, Clone)]
139pub struct Array {
140 pub id: TermId,
142 pub domain: SortId,
144 pub range: SortId,
146}
147
148impl Array {
149 #[must_use]
151 pub fn from_id(id: TermId, domain: SortId, range: SortId) -> Self {
152 Self { id, domain, range }
153 }
154
155 #[must_use]
160 pub fn new_const(ctx: &Z3Context, name: &str, domain: SortId, range: SortId) -> Self {
161 let arr_sort = ctx.tm.borrow_mut().sorts.array(domain, range);
162 let id = build!(ctx, mk_var, name, arr_sort);
163 Self { id, domain, range }
164 }
165
166 #[must_use]
170 pub fn select(ctx: &Z3Context, arr: &Array, idx: TermId) -> TermId {
171 build!(ctx, mk_select, arr.id, idx)
172 }
173
174 #[must_use]
178 pub fn store(ctx: &Z3Context, arr: &Array, idx: TermId, val: TermId) -> Array {
179 let id = build!(ctx, mk_store, arr.id, idx, val);
180 Array {
181 id,
182 domain: arr.domain,
183 range: arr.range,
184 }
185 }
186
187 #[must_use]
189 pub fn eq(ctx: &Z3Context, lhs: &Array, rhs: &Array) -> Bool {
190 let id = build!(ctx, mk_eq, lhs.id, rhs.id);
191 Bool { id }
192 }
193}
194
195impl From<Array> for TermId {
196 fn from(a: Array) -> Self {
197 a.id
198 }
199}
200
201#[derive(Debug, Clone)]
211pub struct FuncDecl {
212 pub name: String,
214 pub domain: Vec<SortId>,
216 pub range: SortId,
218}
219
220impl FuncDecl {
221 #[must_use]
226 pub fn new(_ctx: &Z3Context, name: &str, domain: &[SortId], range: SortId) -> Self {
227 Self {
228 name: name.to_string(),
229 domain: domain.to_vec(),
230 range,
231 }
232 }
233
234 #[must_use]
243 pub fn apply(&self, ctx: &Z3Context, args: &[TermId]) -> TermId {
244 debug_assert_eq!(
245 args.len(),
246 self.domain.len(),
247 "FuncDecl::apply arity mismatch: declared {}, got {}",
248 self.domain.len(),
249 args.len()
250 );
251 let range = self.range;
252 build!(ctx, mk_apply, &self.name, args.iter().copied(), range)
253 }
254}
255
256#[must_use]
265pub fn forall_bool<'a>(
266 ctx: &Z3Context,
267 vars: impl IntoIterator<Item = (&'a str, SortId)>,
268 body: &Bool,
269) -> Bool {
270 let id = ctx.tm.borrow_mut().mk_forall(vars, body.id);
271 Bool { id }
272}
273
274#[must_use]
281pub fn exists_bool<'a>(
282 ctx: &Z3Context,
283 vars: impl IntoIterator<Item = (&'a str, SortId)>,
284 body: &Bool,
285) -> Bool {
286 let id = ctx.tm.borrow_mut().mk_exists(vars, body.id);
287 Bool { id }
288}
289
290pub struct Z3Optimize {
318 ctx_tm: std::rc::Rc<std::cell::RefCell<oxiz_core::ast::TermManager>>,
320 opt: Optimizer,
322 objectives: Vec<(TermId, ObjectiveDir)>,
324 lower_bounds: Vec<Option<String>>,
326 upper_bounds: Vec<Option<String>>,
328 last_result: SatResult,
330}
331
332#[derive(Debug, Clone, Copy)]
334enum ObjectiveDir {
335 Minimize,
336 Maximize,
337}
338
339impl Z3Optimize {
340 #[must_use]
345 pub fn new(ctx: &Z3Context) -> Self {
346 Self {
347 ctx_tm: ctx.tm.clone(),
348 opt: Optimizer::new(),
349 objectives: Vec::new(),
350 lower_bounds: Vec::new(),
351 upper_bounds: Vec::new(),
352 last_result: SatResult::Unknown,
353 }
354 }
355
356 pub fn assert(&mut self, b: &Bool) {
358 self.opt.assert(b.id);
359 }
360
361 pub fn minimize(&mut self, t: TermId) -> usize {
366 let idx = self.objectives.len();
367 self.opt.minimize(t);
368 self.objectives.push((t, ObjectiveDir::Minimize));
369 self.lower_bounds.push(None);
370 self.upper_bounds.push(None);
371 idx
372 }
373
374 pub fn maximize(&mut self, t: TermId) -> usize {
379 let idx = self.objectives.len();
380 self.opt.maximize(t);
381 self.objectives.push((t, ObjectiveDir::Maximize));
382 self.lower_bounds.push(None);
383 self.upper_bounds.push(None);
384 idx
385 }
386
387 pub fn check(&mut self) -> SatResult {
392 let result = self.opt.optimize(&mut self.ctx_tm.borrow_mut());
397
398 match &result {
399 OptimizationResult::Optimal { value, model: _ } => {
400 let tm = self.ctx_tm.borrow();
401 let val_str = Self::term_to_string(&tm, *value);
402 drop(tm);
403 for idx in 0..self.objectives.len() {
404 self.lower_bounds[idx] = Some(val_str.clone());
405 self.upper_bounds[idx] = Some(val_str.clone());
406 }
407 self.last_result = SatResult::Sat;
408 }
409 OptimizationResult::Unsat => {
410 self.last_result = SatResult::Unsat;
411 }
412 _ => {
413 self.last_result = SatResult::Unknown;
414 }
415 }
416
417 self.last_result
418 }
419
420 #[must_use]
423 pub fn get_lower(&self, idx: usize) -> Option<String> {
424 self.lower_bounds.get(idx).and_then(|b| b.clone())
425 }
426
427 #[must_use]
430 pub fn get_upper(&self, idx: usize) -> Option<String> {
431 self.upper_bounds.get(idx).and_then(|b| b.clone())
432 }
433
434 fn term_to_string(tm: &oxiz_core::ast::TermManager, id: TermId) -> String {
436 use oxiz_core::ast::TermKind;
437 if let Some(t) = tm.get(id) {
438 match &t.kind {
439 TermKind::IntConst(n) => return n.to_string(),
440 TermKind::RealConst(r) => return r.to_string(),
441 TermKind::True => return "true".to_string(),
442 TermKind::False => return "false".to_string(),
443 _ => {}
444 }
445 }
446 format!("term#{}", id.0)
447 }
448}
449
450impl Z3Context {
453 #[must_use]
463 pub fn array_sort(&self, domain: SortId, range: SortId) -> SortId {
464 self.tm.borrow_mut().sorts.array(domain, range)
465 }
466}
467
468#[must_use]
478pub fn real_numeral(ctx: &Z3Context, num: i64, den: i64) -> Real {
479 Real::from_frac(ctx, num, den)
480}
481
482#[must_use]
487pub fn int_numeral(ctx: &Z3Context, value: i64) -> Int {
488 Int::from_i64(ctx, value)
489}
490
491#[cfg(test)]
494mod tests {
495 use super::*;
496 use crate::z3_compat::{Z3Config, Z3Context};
497 use num_bigint::BigInt;
498
499 fn ctx() -> Z3Context {
500 Z3Context::new(&Z3Config::new())
501 }
502
503 #[test]
506 fn test_real_gt_smoke() {
507 let ctx = ctx();
508 let a = Real::new_const(&ctx, "a");
509 let b = Real::new_const(&ctx, "b");
510 let _gt = Real::gt(&ctx, &a, &b);
511 }
512
513 #[test]
514 fn test_real_ge_smoke() {
515 let ctx = ctx();
516 let a = Real::new_const(&ctx, "a");
517 let b = Real::new_const(&ctx, "b");
518 let _ge = Real::ge(&ctx, &a, &b);
519 }
520
521 #[test]
522 fn test_real_neg_smoke() {
523 let ctx = ctx();
524 let a = Real::new_const(&ctx, "a");
525 let _neg = Real::neg(&ctx, &a);
526 }
527
528 #[test]
529 fn test_real_div_smoke() {
530 let ctx = ctx();
531 let a = Real::new_const(&ctx, "a");
532 let b = Real::new_const(&ctx, "b");
533 let _div = Real::div(&ctx, &a, &b);
534 }
535
536 #[test]
537 fn test_real_from_i64_smoke() {
538 let ctx = ctx();
539 let _r = Real::from_i64(&ctx, 42);
540 }
541
542 #[test]
545 fn test_ite_bool_smoke() {
546 let ctx = ctx();
547 let c = Bool::new_const(&ctx, "c");
548 let t = Bool::from_bool(&ctx, true);
549 let e = Bool::from_bool(&ctx, false);
550 let _ite = ite_bool(&ctx, &c, &t, &e);
551 }
552
553 #[test]
554 fn test_ite_int_smoke() {
555 let ctx = ctx();
556 let c = Bool::new_const(&ctx, "c");
557 let t = Int::from_i64(&ctx, 1);
558 let e = Int::from_i64(&ctx, 0);
559 let _ite = ite_int(&ctx, &c, &t, &e);
560 }
561
562 #[test]
563 fn test_ite_real_smoke() {
564 let ctx = ctx();
565 let c = Bool::new_const(&ctx, "c");
566 let t = Real::from_i64(&ctx, 1);
567 let e = Real::from_i64(&ctx, 0);
568 let _ite = ite_real(&ctx, &c, &t, &e);
569 }
570
571 #[test]
572 fn test_ite_bv_width() {
573 let ctx = ctx();
574 let c = Bool::new_const(&ctx, "c");
575 let t = BV::from_u64(&ctx, 1, 32);
576 let e = BV::from_u64(&ctx, 0, 32);
577 let ite = ite_bv(&ctx, &c, &t, &e);
578 assert_eq!(ite.width, 32);
579 }
580
581 #[test]
584 fn test_distinct_int_smoke() {
585 let ctx = ctx();
586 let x = Int::from_i64(&ctx, 1);
587 let y = Int::from_i64(&ctx, 2);
588 let z = Int::from_i64(&ctx, 3);
589 let _d = distinct_int(&ctx, &[x, y, z]);
590 }
591
592 #[test]
593 fn test_distinct_real_smoke() {
594 let ctx = ctx();
595 let a = Real::from_i64(&ctx, 1);
596 let b = Real::from_i64(&ctx, 2);
597 let _d = distinct_real(&ctx, &[a, b]);
598 }
599
600 #[test]
601 fn test_distinct_bv_smoke() {
602 let ctx = ctx();
603 let a = BV::from_u64(&ctx, 0, 8);
604 let b = BV::from_u64(&ctx, 1, 8);
605 let _d = distinct_bv(&ctx, &[a, b]);
606 }
607
608 #[test]
611 fn test_array_new_const_smoke() {
612 let ctx = ctx();
613 let dom = ctx.int_sort();
614 let rng = ctx.int_sort();
615 let _a = Array::new_const(&ctx, "arr", dom, rng);
616 }
617
618 #[test]
619 fn test_array_select_store_terms_constructed() {
620 let ctx = ctx();
621 let dom = ctx.int_sort();
622 let rng = ctx.int_sort();
623 let arr = Array::new_const(&ctx, "arr", dom, rng);
624 let idx = Int::from_i64(&ctx, 0);
625 let val = Int::from_i64(&ctx, 42);
626 let arr2 = Array::store(&ctx, &arr, idx.id, val.id);
627 let _selected = Array::select(&ctx, &arr2, idx.id);
628 }
629
630 #[test]
631 fn test_array_eq_smoke() {
632 let ctx = ctx();
633 let dom = ctx.int_sort();
634 let rng = ctx.int_sort();
635 let a = Array::new_const(&ctx, "a", dom, rng);
636 let b = Array::new_const(&ctx, "b", dom, rng);
637 let _eq = Array::eq(&ctx, &a, &b);
638 }
639
640 #[test]
641 fn test_array_sort_helper() {
642 let ctx = ctx();
643 let dom = ctx.int_sort();
644 let rng = ctx.bool_sort();
645 let _sort = ctx.array_sort(dom, rng);
646 }
647
648 #[test]
651 fn test_func_decl_apply_smoke() {
652 let ctx = ctx();
653 let int_sort = ctx.int_sort();
654 let f = FuncDecl::new(&ctx, "f", &[int_sort], int_sort);
655 let arg = Int::new_const(&ctx, "x");
656 let _result = f.apply(&ctx, &[arg.id]);
657 }
658
659 #[test]
660 fn test_func_decl_two_args() {
661 let ctx = ctx();
662 let int_sort = ctx.int_sort();
663 let bool_sort = ctx.bool_sort();
664 let g = FuncDecl::new(&ctx, "g", &[int_sort, int_sort], bool_sort);
665 let x = Int::new_const(&ctx, "x");
666 let y = Int::new_const(&ctx, "y");
667 let _r = g.apply(&ctx, &[x.id, y.id]);
668 }
669
670 #[test]
673 fn test_forall_bool_smoke() {
674 let ctx = ctx();
675 let int_sort = ctx.int_sort();
676 let x_var = ctx.tm.borrow_mut().mk_var("x", int_sort);
678 let zero = ctx.tm.borrow_mut().mk_int(BigInt::from(0));
679 let body_id = ctx.tm.borrow_mut().mk_ge(x_var, zero);
680 let body = Bool { id: body_id };
681 let _q = forall_bool(&ctx, [("x", int_sort)], &body);
683 }
684
685 #[test]
686 fn test_exists_bool_smoke() {
687 let ctx = ctx();
688 let int_sort = ctx.int_sort();
689 let x_var = ctx.tm.borrow_mut().mk_var("x", int_sort);
690 let zero = ctx.tm.borrow_mut().mk_int(BigInt::from(0));
691 let body_id = ctx.tm.borrow_mut().mk_ge(x_var, zero);
692 let body = Bool { id: body_id };
693 let _q = exists_bool(&ctx, [("x", int_sort)], &body);
694 }
695
696 #[test]
699 fn test_optimize_sat_no_objectives() {
700 let cfg = Z3Config::new();
702 let ctx = Z3Context::new(&cfg);
703 let mut opt = Z3Optimize::new(&ctx);
704 let t = Bool::from_bool(&ctx, true);
705 opt.assert(&t);
706 let result = opt.check();
707 assert!(
710 result == SatResult::Sat || result == SatResult::Unknown,
711 "Expected Sat or Unknown, got {:?}",
712 result
713 );
714 }
715
716 #[test]
717 fn test_optimize_minimize_term_constructed() {
718 let cfg = Z3Config::new();
719 let ctx = Z3Context::new(&cfg);
720 let mut opt = Z3Optimize::new(&ctx);
721 let x = Int::new_const(&ctx, "x");
723 let zero = Int::from_i64(&ctx, 0);
724 let ge = Int::ge(&ctx, &x, &zero);
725 opt.assert(&ge);
726 let _idx = opt.minimize(x.id);
727 let _result = opt.check();
729 }
730
731 #[test]
732 fn test_optimize_get_lower_before_check_is_none() {
733 let cfg = Z3Config::new();
734 let ctx = Z3Context::new(&cfg);
735 let mut opt = Z3Optimize::new(&ctx);
736 let x = Int::new_const(&ctx, "x");
737 let _idx = opt.minimize(x.id);
738 assert!(opt.get_lower(0).is_none());
740 assert!(opt.get_upper(0).is_none());
741 }
742}