Implement bound python ops for Bool and BV (#12)

twizmwazin · Oct 31, 2024 · 39247dc · 39247dc
1 parent b706d9b
commit 39247dc
Show file tree

Hide file tree

Showing 12 changed files with 510 additions and 23 deletions.
diff --git a/crates/clarirs_core/src/algorithms/simplify.rs b/crates/clarirs_core/src/algorithms/simplify.rs
@@ -69,6 +69,7 @@ impl<'c> Simplify<'c> for BitVecAst<'c> {
                 BitVecOp::And(arc, arc1) => todo!(),
                 BitVecOp::Or(arc, arc1) => todo!(),
                 BitVecOp::Xor(arc, arc1) => todo!(),
+                BitVecOp::Abs(arc) => todo!(),
                 BitVecOp::Add(arc, arc1) => todo!(),
                 BitVecOp::Sub(arc, arc1) => todo!(),
                 BitVecOp::Mul(arc, arc1) => todo!(),

diff --git a/crates/clarirs_core/src/ast/bitvec.rs b/crates/clarirs_core/src/ast/bitvec.rs
@@ -14,6 +14,7 @@ pub enum BitVecOp<'c> {
     And(BitVecAst<'c>, BitVecAst<'c>),
     Or(BitVecAst<'c>, BitVecAst<'c>),
     Xor(BitVecAst<'c>, BitVecAst<'c>),
+    Abs(BitVecAst<'c>),
     Add(BitVecAst<'c>, BitVecAst<'c>),
     Sub(BitVecAst<'c>, BitVecAst<'c>),
     Mul(BitVecAst<'c>, BitVecAst<'c>),
@@ -50,6 +51,7 @@ impl<'c> Op<'c> for BitVecOp<'c> {
             BitVecOp::BVS(..) | BitVecOp::BVV(..) | BitVecOp::SI(..) => vec![],
 
             BitVecOp::Not(a)
+            | BitVecOp::Abs(a)
             | BitVecOp::Reverse(a)
             | BitVecOp::ZeroExt(a, _)
             | BitVecOp::SignExt(a, _)

diff --git a/crates/clarirs_core/src/ast/factory.rs b/crates/clarirs_core/src/ast/factory.rs
@@ -96,6 +96,13 @@ pub trait AstFactory<'c>: Sized {
         Op::neq(self, lhs, rhs)
     }
 
+    fn abs<Op: SupportsAbs<'c>>(
+        &'c self,
+        ast: &AstRef<'c, Op>,
+    ) -> Result<AstRef<'c, Op>, ClarirsError> {
+        Op::abs(self, ast)
+    }
+
     fn add<Op: SupportsAdd<'c>>(
         &'c self,
         lhs: &AstRef<'c, Op>,

diff --git a/crates/clarirs_core/src/ast/factory_support.rs b/crates/clarirs_core/src/ast/factory_support.rs
@@ -21,6 +21,7 @@ macro_rules! uniop_support_trait {
 }
 
 uniop_support_trait!(Not, BooleanOp<'c>, make_bool, BitVecOp<'c>, make_bitvec);
+uniop_support_trait!(Abs, BitVecOp<'c>, make_bitvec);
 
 macro_rules! binop_support_trait {
     ($name:ident, $($impler:ty, $factory_func:ident),*) => {

diff --git a/crates/clarirs_py/src/ast/args.rs b/crates/clarirs_py/src/ast/args.rs
@@ -68,7 +68,7 @@ impl ExtractPyArgs for BitVecOp<'static> {
             BitVecOp::BVS(name, size) => vec![name.to_object(py), size.to_object(py)],
             BitVecOp::BVV(bit_vec) => vec![bit_vec.as_biguint().to_object(py)],
             BitVecOp::SI(_, bit_vec, bit_vec1, bit_vec2, _) => todo!(),
-            BitVecOp::Not(expr) => vec![BV::new(py, expr)?.into_any()],
+            BitVecOp::Not(expr) | BitVecOp::Abs(expr) => vec![BV::new(py, expr)?.into_any()],
             BitVecOp::And(lhs, rhs)
             | BitVecOp::Or(lhs, rhs)
             | BitVecOp::Xor(lhs, rhs)

diff --git a/crates/clarirs_py/src/ast/bool.rs b/crates/clarirs_py/src/ast/bool.rs
@@ -86,25 +86,51 @@ impl Bool {
         self.inner.depth() == 1
     }
 
-    fn is_true(self_: PyRef<Self>) -> bool {
-        self_.inner.is_true()
+    fn is_true(&self) -> bool {
+        self.inner.is_true()
     }
 
-    fn is_false(self_: PyRef<Self>) -> bool {
-        self_.inner.is_false()
+    fn is_false(&self) -> bool {
+        self.inner.is_false()
+    }
+
+    fn __invert__(&self, py: Python) -> Result<Py<Bool>, ClaripyError> {
+        Bool::new(py, &GLOBAL_CONTEXT.not(&self.inner)?)
+    }
+
+    fn __and__(&self, py: Python, other: CoerceBool) -> Result<Py<Bool>, ClaripyError> {
+        Bool::new(
+            py,
+            &GLOBAL_CONTEXT.and(&self.inner, &<CoerceBool as Into<BoolAst>>::into(other))?,
+        )
+    }
+
+    fn __or__(&self, py: Python, other: CoerceBool) -> Result<Py<Bool>, ClaripyError> {
+        Bool::new(
+            py,
+            &GLOBAL_CONTEXT.or(&self.inner, &<CoerceBool as Into<BoolAst>>::into(other))?,
+        )
+    }
+
+    fn __xor__(&self, py: Python, other: CoerceBool) -> Result<Py<Bool>, ClaripyError> {
+        Bool::new(
+            py,
+            &GLOBAL_CONTEXT.xor(&self.inner, &<CoerceBool as Into<BoolAst>>::into(other))?,
+        )
     }
 }
 
 #[pyfunction]
 pub fn BoolS(py: Python, name: &str) -> Result<Py<Bool>, ClaripyError> {
     Bool::new(py, &GLOBAL_CONTEXT.bools(name)?)
 }
+
 #[pyfunction]
 pub fn BoolV(py: Python, value: bool) -> Result<Py<Bool>, ClaripyError> {
     Bool::new(py, &GLOBAL_CONTEXT.boolv(value)?)
 }
 
-#[pyfunction]
+#[pyfunction(name = "Eq")]
 pub fn Eq_(py: Python, a: Bound<Bool>, b: Bound<Bool>) -> Result<Py<Bool>, ClaripyError> {
     Bool::new(py, &GLOBAL_CONTEXT.eq_(&a.get().inner, &b.get().inner)?)
 }
@@ -114,19 +140,6 @@ pub fn Neq(py: Python, a: Bound<Bool>, b: Bound<Bool>) -> Result<Py<Bool>, Clari
     Bool::new(py, &GLOBAL_CONTEXT.neq(&a.get().inner, &b.get().inner)?)
 }
 
-#[pyfunction]
-pub fn If(
-    py: Python,
-    cond: Bound<Bool>,
-    then_: Bound<Bool>,
-    else_: Bound<Bool>,
-) -> Result<Py<Bool>, ClaripyError> {
-    Bool::new(
-        py,
-        &GLOBAL_CONTEXT.if_(&cond.get().inner, &then_.get().inner, &else_.get().inner)?,
-    )
-}
-
 #[pyfunction(name = "true")]
 pub fn true_op(py: Python) -> Result<Py<Bool>, ClaripyError> {
     Bool::new(py, &GLOBAL_CONTEXT.true_()?)
@@ -139,7 +152,19 @@ pub fn false_op(py: Python) -> Result<Py<Bool>, ClaripyError> {
 pub(crate) fn import(_: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<Bool>()?;
 
-    add_pyfunctions!(m, BoolS, BoolV, Not, And, Or, Xor, Eq_, If, true_op, false_op,);
+    add_pyfunctions!(
+        m,
+        BoolS,
+        BoolV,
+        Not,
+        And,
+        Or,
+        Xor,
+        Eq_,
+        super::If,
+        true_op,
+        false_op,
+    );
 
     Ok(())
 }