Skip to content

Commit

Permalink
Improve BV coercion (#20)
Browse files Browse the repository at this point in the history 
* Improve BV coercion

* cargo fmt
  • Loading branch information
@twizmwazin
twizmwazin authored Nov 1, 2024
1 parent 3a31d0a commit c1ff5eb
Show file tree
Hide file tree
Showing 4 changed files with 128 additions and 92 deletions.
2 changes: 1 addition & 1 deletion crates/clarirs_py/src/ast/args.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,9 +95,9 @@ impl ExtractPyArgs for BitVecOp<'static> {
vec![BV::new(py, expr)?.into_any(), amount.to_object(py)]
}
BitVecOp::Extract(expr, end, start) => vec![
BV::new(py, expr)?.into_any(),
end.to_object(py),
start.to_object(py),
BV::new(py, expr)?.into_any(),
],
BitVecOp::Reverse(expr) => vec![BV::new(py, expr)?.into_any()],
BitVecOp::FpToIEEEBV(expr) => vec![FP::new(py, expr)?.into_any()],
Expand Down
136 changes: 60 additions & 76 deletions crates/clarirs_py/src/ast/bv.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ impl BV {
pub fn __add__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.add(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.add(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -148,7 +148,7 @@ impl BV {
pub fn __sub__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.sub(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.sub(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -159,7 +159,7 @@ impl BV {
pub fn __mul__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.mul(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.mul(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -170,7 +170,7 @@ impl BV {
pub fn __truediv__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.udiv(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.udiv(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -181,7 +181,7 @@ impl BV {
pub fn __floordiv__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.udiv(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.udiv(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -198,7 +198,7 @@ impl BV {
// TODO: handle modulo
BV::new(
py,
&GLOBAL_CONTEXT.pow(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.pow(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -214,7 +214,7 @@ impl BV {
pub fn __mod__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.urem(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.urem(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -225,21 +225,21 @@ impl BV {
pub fn SDiv(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.sdiv(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.sdiv(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn SMod(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.srem(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.srem(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __and__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.and(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.and(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -250,7 +250,7 @@ impl BV {
pub fn __or__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.or(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.or(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -261,7 +261,7 @@ impl BV {
pub fn __xor__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.xor(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.xor(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -272,7 +272,7 @@ impl BV {
pub fn __lshift__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.shl(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.shl(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -283,7 +283,7 @@ impl BV {
pub fn __rshift__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.ashr(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ashr(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -294,7 +294,7 @@ impl BV {
pub fn LShR(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.lshr(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.lshr(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -317,98 +317,98 @@ impl BV {
pub fn __eq__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.eq_(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.eq_(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __ne__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.neq(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.neq(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __lt__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ult(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ult(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __le__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ule(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ule(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __gt__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ugt(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ugt(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn __ge__(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.uge(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.uge(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn ULT(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ult(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ult(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn ULE(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ule(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ule(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn UGT(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.ugt(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.ugt(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn UGE(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.uge(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.uge(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn SLT(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.slt(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.slt(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn SLE(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.sle(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.sle(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn SGT(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.sgt(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.sgt(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

pub fn SGE(&self, py: Python, other: CoerceBV) -> Result<Py<Bool>, ClaripyError> {
Bool::new(
py,
&GLOBAL_CONTEXT.sge(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.sge(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand All @@ -427,7 +427,7 @@ impl BV {
pub fn concat(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.concat(&self.inner, &<CoerceBV as Into<BitVecAst>>::into(other))?,
&GLOBAL_CONTEXT.concat(&self.inner, &other.extract_like(py, self).get().inner)?,
)
}

Expand Down Expand Up @@ -509,35 +509,32 @@ pub fn BVV(py: Python, value: Bound<PyAny>, size: Option<u32>) -> Result<Py<BV>,
}

macro_rules! binop {
($name:ident, $context_method:ident, $return_type:ty, $lhs:ty, $rhs:ty) => {
($name:ident, $context_method:ident, $ret:ty) => {
#[pyfunction]
pub fn $name(
py: Python,
lhs: Bound<$lhs>,
rhs: Bound<$rhs>,
) -> Result<Py<$return_type>, ClaripyError> {
<$return_type>::new(
pub fn $name(py: Python, lhs: CoerceBV, rhs: CoerceBV) -> Result<Py<$ret>, ClaripyError> {
let (elhs, erhs) = CoerceBV::extract_pair(py, &lhs, &rhs);
<$ret>::new(
py,
&GLOBAL_CONTEXT.$context_method(&lhs.get().inner, &rhs.get().inner)?,
&GLOBAL_CONTEXT.$context_method(&elhs.get().inner, &erhs.get().inner)?,
)
}
};
}

binop!(Add, add, BV, BV, BV);
binop!(Sub, sub, BV, BV, BV);
binop!(Mul, mul, BV, BV, BV);
binop!(UDiv, udiv, BV, BV, BV);
binop!(SDiv, sdiv, BV, BV, BV);
binop!(UMod, urem, BV, BV, BV);
binop!(SMod, srem, BV, BV, BV);
binop!(Pow, pow, BV, BV, BV);
binop!(ShL, shl, BV, BV, BV);
binop!(LShR, lshr, BV, BV, BV);
binop!(AShR, ashr, BV, BV, BV);
binop!(RotateLeft, rotate_left, BV, BV, BV);
binop!(RotateRight, rotate_right, BV, BV, BV);
binop!(Concat, concat, BV, BV, BV);
binop!(Add, add, BV);
binop!(Sub, sub, BV);
binop!(Mul, mul, BV);
binop!(UDiv, udiv, BV);
binop!(SDiv, sdiv, BV);
binop!(UMod, urem, BV);
binop!(SMod, srem, BV);
binop!(Pow, pow, BV);
binop!(ShL, shl, BV);
binop!(LShR, lshr, BV);
binop!(AShR, ashr, BV);
binop!(RotateLeft, rotate_left, BV);
binop!(RotateRight, rotate_right, BV);
binop!(Concat, concat, BV);

#[pyfunction]
pub fn Extract(py: Python, base: Bound<BV>, start: u32, end: u32) -> Result<Py<BV>, ClaripyError> {
Expand All @@ -559,29 +556,16 @@ pub fn Reverse(py: Python, base: Bound<BV>) -> Result<Py<BV>, ClaripyError> {
BV::new(py, &GLOBAL_CONTEXT.reverse(&base.get().inner)?)
}

binop!(ULT, ult, Bool, BV, BV);
binop!(ULE, ule, Bool, BV, BV);
binop!(UGT, ugt, Bool, BV, BV);
binop!(UGE, uge, Bool, BV, BV);
binop!(SLT, slt, Bool, BV, BV);
binop!(SLE, sle, Bool, BV, BV);
binop!(SGT, sgt, Bool, BV, BV);
binop!(SGE, sge, Bool, BV, BV);
binop!(Eq_, eq_, Bool, BV, BV);
binop!(Neq, neq, Bool, BV, BV);

#[pyfunction]
pub fn If(
py: Python,
cond: Bound<Bool>,
then_: Bound<BV>,
else_: Bound<BV>,
) -> Result<Py<BV>, ClaripyError> {
BV::new(
py,
&GLOBAL_CONTEXT.if_(&cond.get().inner, &then_.get().inner, &else_.get().inner)?,
)
}
binop!(ULT, ult, Bool);
binop!(ULE, ule, Bool);
binop!(UGT, ugt, Bool);
binop!(UGE, uge, Bool);
binop!(SLT, slt, Bool);
binop!(SLE, sle, Bool);
binop!(SGT, sgt, Bool);
binop!(SGE, sge, Bool);
binop!(Eq_, eq_, Bool);
binop!(Neq, neq, Bool);

pub(crate) fn import(_: Python, m: &Bound<PyModule>) -> PyResult<()> {
m.add_class::<BV>()?;
Expand Down Expand Up @@ -621,7 +605,7 @@ pub(crate) fn import(_: Python, m: &Bound<PyModule>) -> PyResult<()> {
SGT,
SGE,
Eq_,
If,
super::If,
);

Ok(())
Expand Down
Loading

0 comments on commit c1ff5eb

Please sign in to comment.