1#![doc = concat!("```ignore\n", include_str!("./FixedSizeBitVectors.smt2"), "```")]
2use itertools::Itertools;
5use smtlib_lowlevel::{
6 ast::{self, Identifier, Index, QualIdentifier, Term},
7 lexicon::{self, Numeral},
8 Storage,
9};
10
11use crate::{
12 sorts::Sort,
13 terms::{
14 app, qual_ident, ApplicationArgs, Const, Dynamic, IntoWithStorage, STerm, Sorted,
15 StaticSorted,
16 },
17 Bool,
18};
19
20#[derive(Debug, Clone, Copy)]
25pub struct BitVec<'st, const M: usize>(STerm<'st>);
26impl<'st, const M: usize> From<Const<'st, BitVec<'st, M>>> for BitVec<'st, M> {
27 fn from(c: Const<'st, BitVec<'st, M>>) -> Self {
28 c.1
29 }
30}
31impl<const M: usize> std::fmt::Display for BitVec<'_, M> {
32 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 self.term().fmt(f)
34 }
35}
36
37impl<'st, const M: usize> From<BitVec<'st, M>> for Dynamic<'st> {
38 fn from(i: BitVec<'st, M>) -> Self {
39 i.into_dynamic()
40 }
41}
42
43impl<'st, const M: usize> From<BitVec<'st, M>> for STerm<'st> {
44 fn from(i: BitVec<'st, M>) -> Self {
45 i.0
46 }
47}
48impl<'st, const M: usize> From<STerm<'st>> for BitVec<'st, M> {
49 fn from(t: STerm<'st>) -> Self {
50 BitVec(t)
51 }
52}
53
54impl<'st, const M: usize> From<(STerm<'st>, Sort<'st>)> for BitVec<'st, M> {
55 fn from((t, _): (STerm<'st>, Sort<'st>)) -> Self {
56 t.into()
57 }
58}
59
60fn i64_to_bit_array<const M: usize>(i: i64) -> [bool; M] {
61 std::array::from_fn(|idx| (i >> (M - idx - 1)) & 1 == 1)
62}
63
64impl<'st, const M: usize> TryFrom<BitVec<'st, M>> for i64 {
70 type Error = std::num::ParseIntError;
71
72 fn try_from(value: BitVec<'st, M>) -> Result<Self, Self::Error> {
73 match value.term() {
74 Term::SpecConstant(c) => match c {
75 ast::SpecConstant::Numeral(_) => todo!(),
76 ast::SpecConstant::Decimal(_) => todo!(),
77 ast::SpecConstant::Hexadecimal(h) => h.parse(),
78 ast::SpecConstant::Binary(b) => b.parse(),
79 ast::SpecConstant::String(_) => todo!(),
80 },
81 _ => todo!(),
82 }
83 }
84}
85impl<'st, const M: usize> TryFrom<BitVec<'st, M>> for [bool; M] {
86 type Error = std::num::ParseIntError;
87
88 fn try_from(value: BitVec<'st, M>) -> Result<Self, Self::Error> {
89 Ok(i64_to_bit_array(value.try_into()?))
90 }
91}
92
93impl<'st, const M: usize> StaticSorted<'st> for BitVec<'st, M> {
94 type Inner = Self;
95 const AST_SORT: ast::Sort<'static> =
96 ast::Sort::new_indexed("BitVec", &[Index::Numeral(lexicon::Numeral::from_usize(M))]);
97 fn static_st(&self) -> &'st Storage {
98 self.sterm().st()
99 }
100
101 fn sort() -> Sort<'st> {
102 Self::AST_SORT.into()
103 }
104
105 fn new_const(st: &'st Storage, name: &str) -> Const<'st, Self> {
106 let name = st.alloc_str(name);
107 let bv = Term::Identifier(qual_ident(name, Some(Self::AST_SORT)));
108 let bv = STerm::new(st, bv);
109 Const(name, bv.into())
110 }
111}
112impl<'st, const M: usize> IntoWithStorage<'st, BitVec<'st, M>> for [bool; M] {
113 fn into_with_storage(self, st: &'st Storage) -> BitVec<'st, M> {
114 let term = Term::Identifier(qual_ident(
115 st.alloc_str(&format!("#b{}", self.iter().map(|x| *x as u8).format(""))),
116 None,
117 ));
118 STerm::new(st, term).into()
119 }
120}
121impl<'st, const M: usize> IntoWithStorage<'st, BitVec<'st, M>> for i64 {
122 fn into_with_storage(self, st: &'st Storage) -> BitVec<'st, M> {
123 i64_to_bit_array(self).into_with_storage(st)
124 }
125}
126
127pub trait PrimInt<const M: usize>: Sized {
129 fn to_bit_array(&self) -> [bool; M];
131 fn from_spec_constant<'st>(
133 spec_constant: &ast::SpecConstant<'st>,
134 ) -> Result<Self, std::num::ParseIntError>;
135}
136
137impl PrimInt<8> for u8 {
138 fn to_bit_array(&self) -> [bool; 8] {
139 std::array::from_fn(|idx| (self >> (8 - idx - 1)) & 1 == 1)
140 }
141 fn from_spec_constant<'st>(
142 spec_constant: &ast::SpecConstant<'st>,
143 ) -> Result<Self, std::num::ParseIntError> {
144 match spec_constant {
145 ast::SpecConstant::Hexadecimal(h) => Self::from_str_radix(&h.0[2..], 16),
146 ast::SpecConstant::Binary(b) => Self::from_str_radix(&b.0[2..], 2),
147 _ => panic!(),
148 }
149 }
150}
151impl PrimInt<32> for u32 {
152 fn to_bit_array(&self) -> [bool; 32] {
153 std::array::from_fn(|idx| (self >> (32 - idx - 1)) & 1 == 1)
154 }
155 fn from_spec_constant<'st>(
156 spec_constant: &ast::SpecConstant<'st>,
157 ) -> Result<Self, std::num::ParseIntError> {
158 match spec_constant {
159 ast::SpecConstant::Hexadecimal(h) => Self::from_str_radix(&h.0[2..], 16),
160 ast::SpecConstant::Binary(b) => Self::from_str_radix(&b.0[2..], 2),
161 _ => panic!(),
162 }
163 }
164}
165impl PrimInt<64> for u64 {
166 fn to_bit_array(&self) -> [bool; 64] {
167 std::array::from_fn(|idx| (self >> (64 - idx - 1)) & 1 == 1)
168 }
169 fn from_spec_constant<'st>(
170 spec_constant: &ast::SpecConstant<'st>,
171 ) -> Result<Self, std::num::ParseIntError> {
172 match spec_constant {
173 ast::SpecConstant::Hexadecimal(h) => Self::from_str_radix(&h.0[2..], 16),
174 ast::SpecConstant::Binary(b) => Self::from_str_radix(&b.0[2..], 2),
175 _ => panic!(),
176 }
177 }
178}
179
180impl<'st, const M: usize> BitVec<'st, M> {
181 pub fn new(
187 st: &'st Storage,
188 value: impl IntoWithStorage<'st, BitVec<'st, M>>,
189 ) -> BitVec<'st, M> {
190 value.into_with_storage(st)
191 }
192 pub fn new_prim(st: &'st Storage, value: impl PrimInt<M>) -> BitVec<'st, M> {
195 Self::new(st, value.to_bit_array())
196 }
197 pub fn try_into_prim<T: PrimInt<M>>(&self) -> Result<T, std::num::ParseIntError> {
200 match self.term() {
201 Term::SpecConstant(spec_constant) => T::from_spec_constant(spec_constant),
202 _ => panic!(),
203 }
204 }
205 fn binop<T: From<STerm<'st>>>(self, op: &'st str, other: BitVec<'st, M>) -> T {
207 app(self.st(), op, (self.term(), other.term())).into()
208 }
209 fn unop<T: From<STerm<'st>>>(self, op: &'st str) -> T {
210 app(self.st(), op, self.term()).into()
211 }
212 fn unop_indexed<T: From<STerm<'st>>>(
213 self,
214 op: &'st str,
215 index: impl IntoIterator<Item = usize>,
216 ) -> T {
217 let index = index
218 .into_iter()
219 .map(|i| Index::Numeral(Numeral::from_usize(i)))
220 .collect_vec();
221 let qual_id =
222 QualIdentifier::Identifier(Identifier::indexed(op, self.st().alloc_slice(&index)));
223
224 STerm::new(
225 self.st(),
226 Term::Application(qual_id, self.term().into_args(self.st())),
227 )
228 .into()
229 }
230
231 #[cfg(feature = "const-bit-vec")]
232 pub fn extract<const I: usize, const J: usize>(self) -> BitVec<'st, { I - J + 1 }> {
240 assert!(M > I);
241 assert!(I >= J);
242
243 self.extract_(I, J)
244 }
245
246 pub fn extract_<const O: usize>(self, i: usize, j: usize) -> BitVec<'st, O> {
260 assert!(M > i);
261 assert!(i >= j);
262 assert_eq!(O, i - j + 1);
263
264 self.unop_indexed("extract", [i, j])
265 }
266
267 #[cfg(feature = "const-bit-vec")]
268 pub fn repeat<const N: usize>(self) -> BitVec<'st, { N * M }> {
270 self.repeat_(N)
271 }
272
273 pub fn repeat_<const O: usize>(self, count: usize) -> BitVec<'st, O> {
280 assert!(count > 0);
281 assert_eq!(O, M * count);
282 self.unop_indexed("repeat", [count])
283 }
284
285 #[cfg(feature = "const-bit-vec")]
286 pub fn concat<const N: usize>(
289 self,
290 other: impl Into<BitVec<'st, N>>,
291 ) -> BitVec<'st, { N + M }> {
292 self.concat_(other)
293 }
294
295 pub fn concat_<const N: usize, const O: usize>(
303 self,
304 other: impl Into<BitVec<'st, N>>,
305 ) -> BitVec<'st, O> {
306 assert_eq!(O, N + M);
307 app(self.st(), "concat", (self.term(), other.into().term())).into()
308 }
309
310 #[cfg(feature = "const-bit-vec")]
311 pub fn zero_extend<const I: usize>(self) -> BitVec<'st, { M + I }> {
313 self.zero_extend_(I)
314 }
315
316 pub fn zero_extend_<const O: usize>(self, n: usize) -> BitVec<'st, O> {
323 assert_eq!(M + n, O);
324 self.unop_indexed("zero_extend", [n])
325 }
326
327 #[cfg(feature = "const-bit-vec")]
328 pub fn sign_extend<const I: usize>(self) -> BitVec<'st, { M + I }> {
330 self.sign_extend_(I)
331 }
332
333 pub fn sign_extend_<const O: usize>(self, n: usize) -> BitVec<'st, O> {
340 assert_eq!(M + n, O);
341 self.unop_indexed("sign_extend", [n])
342 }
343
344 pub fn rotate_left(self, n: usize) -> Self {
348 self.unop_indexed("rotate_left", [n])
349 }
350 pub fn rotate_right(self, n: usize) -> Self {
352 self.unop_indexed("rotate_right", [n])
353 }
354
355 pub fn bvnot(self) -> Self {
358 self.unop("bvnot")
359 }
360 pub fn bvneg(self) -> Self {
362 self.unop("bvneg")
363 }
364
365 pub fn bvnand(self, other: impl Into<Self>) -> Self {
368 self.binop("bvnand", other.into())
369 }
370 pub fn bvnor(self, other: impl Into<Self>) -> Self {
372 self.binop("bvnor", other.into())
373 }
374 pub fn bvxnor(self, other: impl Into<Self>) -> Self {
376 self.binop("bvxnor", other.into())
377 }
378 pub fn bvcomp(self, other: impl Into<Self>) -> BitVec<'st, 1> {
380 self.binop("bvcomp", other.into())
381 }
382 pub fn bvsdiv(self, other: impl Into<Self>) -> Self {
384 self.binop("bvsdiv", other.into())
385 }
386 pub fn bvsrem(self, other: impl Into<Self>) -> Self {
389 self.binop("bvsrem", other.into())
390 }
391 pub fn bvsmod(self, other: impl Into<Self>) -> Self {
394 self.binop("bvsmod", other.into())
395 }
396 pub fn bvashr(self, other: impl Into<Self>) -> Self {
398 self.binop("bvashr", other.into())
399 }
400 pub fn bvult(self, other: impl Into<Self>) -> Bool<'st> {
402 self.binop("bvult", other.into())
403 }
404 pub fn bvule(self, other: impl Into<Self>) -> Bool<'st> {
406 self.binop("bvule", other.into())
407 }
408 pub fn bvugt(self, other: impl Into<Self>) -> Bool<'st> {
410 self.binop("bvugt", other.into())
411 }
412 pub fn bvuge(self, other: impl Into<Self>) -> Bool<'st> {
414 self.binop("bvuge", other.into())
415 }
416 pub fn bvslt(self, other: impl Into<Self>) -> Bool<'st> {
418 self.binop("bvslt", other.into())
419 }
420 pub fn bvsle(self, other: impl Into<Self>) -> Bool<'st> {
422 self.binop("bvsle", other.into())
423 }
424 pub fn bvsgt(self, other: impl Into<Self>) -> Bool<'st> {
426 self.binop("bvsgt", other.into())
427 }
428 pub fn bvsge(self, other: impl Into<Self>) -> Bool<'st> {
430 self.binop("bvsge", other.into())
431 }
432}
433
434impl<const M: usize> std::ops::Not for BitVec<'_, M> {
435 type Output = Self;
436 fn not(self) -> Self::Output {
437 self.bvnot()
438 }
439}
440impl<const M: usize> std::ops::Neg for BitVec<'_, M> {
441 type Output = Self;
442 fn neg(self) -> Self::Output {
443 self.bvneg()
444 }
445}
446
447macro_rules! impl_op {
448 ($ty:ty, $other:ty, $trait:tt, $fn:ident, $op:ident, $a_trait:tt, $a_fn:tt, $a_op:tt) => {
449 impl<'st, const M: usize, R> std::ops::$trait<R> for Const<'st, $ty>
450 where
451 R: Into<$ty>,
452 {
453 type Output = $ty;
454 fn $fn(self, rhs: R) -> Self::Output {
455 self.1.binop(stringify!($op), rhs.into())
456 }
457 }
458 impl<'st, const M: usize, R> std::ops::$trait<R> for $ty
459 where
460 R: Into<$ty>,
461 {
462 type Output = Self;
463 fn $fn(self, rhs: R) -> Self::Output {
464 self.binop(stringify!($op), rhs.into())
465 }
466 }
467 impl<'st, const M: usize, R> std::ops::$a_trait<R> for $ty
480 where
481 R: Into<$ty>,
482 {
483 fn $a_fn(&mut self, rhs: R) {
484 *self = *self $a_op rhs;
485 }
486 }
487 impl<'st, const M: usize> $ty {
488 #[doc = concat!("Calls `(", stringify!($op), " self other)`")]
489 pub fn $op(self, other: impl Into<Self>) -> Self {
490 self.binop(stringify!($op), other.into())
491 }
492 }
493 };
494}
495
496impl_op!(BitVec<'st, M>, [bool; M], BitAnd, bitand, bvand, BitAndAssign, bitand_assign, &);
497impl_op!(BitVec<'st, M>, [bool; M], BitOr, bitor, bvor, BitOrAssign, bitor_assign, |);
498impl_op!(BitVec<'st, M>, [bool; M], BitXor, bitxor, bvxor, BitXorAssign, bitxor_assign, ^);
499impl_op!(BitVec<'st, M>, [bool; M], Add, add, bvadd, AddAssign, add_assign, +);
500impl_op!(BitVec<'st, M>, [bool; M], Sub, sub, bvsub, SubAssign, sub_assign, -);
501impl_op!(BitVec<'st, M>, [bool; M], Mul, mul, bvmul, MulAssign, mul_assign, *);
502impl_op!(BitVec<'st, M>, [bool; M], Div, div, bvudiv, DivAssign, div_assign, /);
503impl_op!(BitVec<'st, M>, [bool; M], Rem, rem, bvurem, RemAssign, rem_assign, %);
504impl_op!(BitVec<'st, M>, [bool; M], Shr, shr, bvlshr, ShrAssign, shr_assign, >>);
505impl_op!(BitVec<'st, M>, [bool; M], Shl, shl, bvshl, ShlAssign, shl_assign, <<);
506
507#[cfg(feature = "const-bit-vec")]
508#[cfg(test)]
509mod tests {
510 use smtlib_lowlevel::{backend::z3_binary::Z3Binary, Storage};
511
512 use super::BitVec;
513 use crate::{
514 terms::{Sorted, StaticSorted},
515 SatResult, Solver,
516 };
517
518 #[test]
519 fn bit_vec_extract_concat() -> Result<(), Box<dyn std::error::Error>> {
520 let st = Storage::new();
521 let a = BitVec::<6>::new_const(&st, "a");
522 let b = BitVec::<4>::new_const(&st, "b");
523 let c = BitVec::<10>::new_const(&st, "c");
524 let d = BitVec::<6>::new(&st, [true, false, true, true, false, true]);
525
526 let mut solver = Solver::new(&st, Z3Binary::new("z3")?)?;
527
528 solver.assert(a._eq(!d))?;
529 solver.assert(c._eq(a.concat(b)))?;
530 solver.assert(b._eq(a.extract::<5, 2>()))?;
531
532 let model = solver.check_sat_with_model()?.expect_sat()?;
533
534 let a: [bool; 6] = model.eval(a).unwrap().try_into()?;
535 let b: [bool; 4] = model.eval(b).unwrap().try_into()?;
536 let c: [bool; 10] = model.eval(c).unwrap().try_into()?;
537 insta::assert_ron_snapshot!(a, @"(false, true, false, false, true, false)");
538 insta::assert_ron_snapshot!(b, @"(false, true, false, false)");
539 insta::assert_ron_snapshot!(c, @"(false, true, false, false, true, false, false, true, false, false)");
540
541 Ok(())
542 }
543
544 #[test]
545 fn bit_vec_math() -> Result<(), Box<dyn std::error::Error>> {
546 let st = Storage::new();
547 let a = BitVec::<6>::new_const(&st, "a");
548 let b = BitVec::<6>::new_const(&st, "b");
549 let c = BitVec::<6>::new_const(&st, "c");
550
551 let mut solver = Solver::new(&st, Z3Binary::new("z3")?)?;
552
553 solver.assert(a._eq(10))?;
554 solver.assert(b._eq(3))?;
555 solver.assert(c._eq(a + b))?;
557
558 assert_eq!(solver.check_sat()?, SatResult::Sat);
559 let model = solver.get_model()?;
560
561 let a: i64 = model.eval(a).unwrap().try_into()?;
562 let b: i64 = model.eval(b).unwrap().try_into()?;
563 let c: i64 = model.eval(c).unwrap().try_into()?;
564 insta::assert_ron_snapshot!(a, @"10");
565 insta::assert_ron_snapshot!(b, @"3");
566 insta::assert_ron_snapshot!(c, @"13");
567
568 Ok(())
569 }
570
571 #[test]
572 fn bit_vec_dynamic() -> Result<(), Box<dyn std::error::Error>> {
573 let st = Storage::new();
574 let a = BitVec::<64>::new_const(&st, "a");
575
576 assert!(BitVec::<64>::try_from_dynamic(a.into_dynamic()).is_some());
577 assert!(BitVec::<32>::try_from_dynamic(a.into_dynamic()).is_none());
578
579 Ok(())
580 }
581}