Add size() and __len__() to BV and FP (#17)

twizmwazin · Nov 1, 2024 · 23099a4 · 23099a4
1 parent bdf2b74
commit 23099a4
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 0 deletions.
diff --git a/crates/clarirs_core/src/ast/bitvec.rs b/crates/clarirs_core/src/ast/bitvec.rs
@@ -5,6 +5,8 @@ use serde::Serialize;
 
 use crate::prelude::*;
 
+use super::float::FloatExt;
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
 pub enum BitVecOp<'c> {
     BVS(String, u32),
@@ -110,3 +112,45 @@ impl<'c> Op<'c> for BitVecOp<'c> {
         }
     }
 }
+
+pub trait BitVecExt<'c> {
+    fn size(&self) -> u32;
+}
+
+impl<'c> BitVecExt<'c> for BitVecAst<'c> {
+    fn size(&self) -> u32 {
+        match self.op() {
+            BitVecOp::BVS(_, size)
+            | BitVecOp::SI(_, _, _, _, size)
+            | BitVecOp::ZeroExt(_, size)
+            | BitVecOp::SignExt(_, size) => *size,
+            BitVecOp::BVV(bv) => bv.len() as u32,
+            BitVecOp::Not(a)
+            | BitVecOp::Abs(a)
+            | BitVecOp::Reverse(a)
+            | BitVecOp::If(_, a, _)
+            | BitVecOp::Annotated(a, _) => a.size(),
+            BitVecOp::And(a, _)
+            | BitVecOp::Or(a, _)
+            | BitVecOp::Xor(a, _)
+            | BitVecOp::Add(a, _)
+            | BitVecOp::Sub(a, _)
+            | BitVecOp::Mul(a, _)
+            | BitVecOp::UDiv(a, _)
+            | BitVecOp::SDiv(a, _)
+            | BitVecOp::URem(a, _)
+            | BitVecOp::SRem(a, _)
+            | BitVecOp::Pow(a, _)
+            | BitVecOp::ShL(a, _)
+            | BitVecOp::LShR(a, _)
+            | BitVecOp::AShR(a, _)
+            | BitVecOp::RotateLeft(a, _)
+            | BitVecOp::RotateRight(a, _) => a.size(),
+            BitVecOp::Extract(_, high, low) => high - low + 1,
+            BitVecOp::Concat(a, b) => a.size() + b.size(),
+            BitVecOp::FpToIEEEBV(fp) => fp.size(),
+            BitVecOp::FpToUBV(_, _, _) | BitVecOp::FpToSBV(_, _, _) => 64,
+            BitVecOp::StrLen(_) | BitVecOp::StrToBV(_) | BitVecOp::StrIndexOf(_, _, _) => 64,
+        }
+    }
+}
diff --git a/crates/clarirs_core/src/ast/float.rs b/crates/clarirs_core/src/ast/float.rs
@@ -68,3 +68,26 @@ impl<'c> Op<'c> for FloatOp<'c> {
         }
     }
 }
+
+pub trait FloatExt<'c> {
+    fn size(&self) -> u32;
+}
+
+impl<'c> FloatExt<'c> for FloatAst<'c> {
+    fn size(&self) -> u32 {
+        match self.op() {
+            FloatOp::FPS(_, sort) => sort.size(),
+            FloatOp::FPV(value) => value.fsort().size(),
+            FloatOp::FpNeg(a, _)
+            | FloatOp::FpAbs(a, _)
+            | FloatOp::FpSqrt(a, _)
+            | FloatOp::FpAdd(a, _, _)
+            | FloatOp::FpSub(a, _, _)
+            | FloatOp::FpMul(a, _, _)
+            | FloatOp::FpDiv(a, _, _)
+            | FloatOp::If(_, a, _)
+            | FloatOp::Annotated(a, _) => a.size(),
+            FloatOp::FpToFp(_, sort, _) | FloatOp::BvToFpUnsigned(_, sort, _) => sort.size(),
+        }
+    }
+}
diff --git a/crates/clarirs_py/src/ast/bv.rs b/crates/clarirs_py/src/ast/bv.rs
@@ -3,6 +3,7 @@
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::LazyLock;
 
+use clarirs_core::ast::bitvec::BitVecExt;
 use dashmap::DashMap;
 use num_bigint::BigUint;
 use pyo3::exceptions::{PyTypeError, PyValueError};
@@ -108,6 +109,14 @@ impl BV {
         self.inner.depth() == 1
     }
 
+    fn size(&self) -> usize {
+        self.inner.size() as usize
+    }
+
+    fn __len__(&self) -> usize {
+        self.size()
+    }
+
     fn annotate(&self, py: Python, annotation: Bound<PyAny>) -> Result<Py<BV>, ClaripyError> {
         let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
         let annotation_bytes = pickle_dumps

diff --git a/crates/clarirs_py/src/ast/fp.rs b/crates/clarirs_py/src/ast/fp.rs
@@ -5,6 +5,7 @@ use std::sync::{
     LazyLock,
 };
 
+use clarirs_core::ast::float::FloatExt;
 use dashmap::DashMap;
 use pyo3::types::{PyBytes, PyFrozenSet, PyWeakrefReference};
 
@@ -191,6 +192,14 @@ impl FP {
         self.inner.depth() == 1
     }
 
+    fn size(&self) -> usize {
+        self.inner.size() as usize
+    }
+
+    fn __len__(&self) -> usize {
+        self.size()
+    }
+
     fn annotate(&self, py: Python, annotation: Bound<PyAny>) -> Result<Py<FP>, ClaripyError> {
         let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
         let annotation_bytes = pickle_dumps