diff --git a/crates/clarirs_core/src/ast/bitvec.rs b/crates/clarirs_core/src/ast/bitvec.rs index ef58b10..82e47ca 100644 --- a/crates/clarirs_core/src/ast/bitvec.rs +++ b/crates/clarirs_core/src/ast/bitvec.rs @@ -97,4 +97,12 @@ impl<'c> Op<'c> for BitVecOp<'c> { .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect()) } } + + fn get_annotations(&self) -> Vec { + if let BitVecOp::Annotated(inner, anno) = self { + inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + } else { + vec![] + } + } } diff --git a/crates/clarirs_core/src/ast/bool.rs b/crates/clarirs_core/src/ast/bool.rs index c3f859e..d41984f 100644 --- a/crates/clarirs_core/src/ast/bool.rs +++ b/crates/clarirs_core/src/ast/bool.rs @@ -112,4 +112,12 @@ impl<'c> Op<'c> for BooleanOp<'c> { .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect()) } } + + fn get_annotations(&self) -> Vec { + if let BooleanOp::Annotated(inner, anno) = self { + inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + } else { + vec![] + } + } } diff --git a/crates/clarirs_core/src/ast/float.rs b/crates/clarirs_core/src/ast/float.rs index 6a98b82..b6e94ce 100644 --- a/crates/clarirs_core/src/ast/float.rs +++ b/crates/clarirs_core/src/ast/float.rs @@ -55,4 +55,12 @@ impl<'c> Op<'c> for FloatOp<'c> { .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect()) } } + + fn get_annotations(&self) -> Vec { + if let FloatOp::Annotated(inner, anno) = self { + inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + } else { + vec![] + } + } } diff --git a/crates/clarirs_core/src/ast/node.rs b/crates/clarirs_core/src/ast/node.rs index 8dbe76b..331e3b8 100644 --- a/crates/clarirs_core/src/ast/node.rs +++ b/crates/clarirs_core/src/ast/node.rs @@ -116,6 +116,10 @@ impl<'c, O: Op<'c>> Op<'c> for AstNode<'c, O> { fn variables(&self) -> HashSet { self.variables.clone() } + + fn get_annotations(&self) -> Vec { + self.op().get_annotations() + } } pub type AstRef<'c, Op> = Arc>; @@ -161,6 +165,15 @@ impl<'c> Op<'c> for VarAst<'c> { } .clone() } + + fn get_annotations(&self) -> Vec { + match self { + VarAst::Boolean(ast) => ast.get_annotations(), + VarAst::BitVec(ast) => ast.get_annotations(), + VarAst::Float(ast) => ast.get_annotations(), + VarAst::String(ast) => ast.get_annotations(), + } + } } impl<'c> From<&BoolAst<'c>> for VarAst<'c> { diff --git a/crates/clarirs_core/src/ast/op.rs b/crates/clarirs_core/src/ast/op.rs index 6484485..c7b6840 100644 --- a/crates/clarirs_core/src/ast/op.rs +++ b/crates/clarirs_core/src/ast/op.rs @@ -31,4 +31,6 @@ pub trait Op<'c>: Debug + Hash + Serialize { fn concrete(&self) -> bool { self.variables().is_empty() } + + fn get_annotations(&self) -> Vec; } diff --git a/crates/clarirs_core/src/ast/string.rs b/crates/clarirs_core/src/ast/string.rs index a1cb2d5..26a601e 100644 --- a/crates/clarirs_core/src/ast/string.rs +++ b/crates/clarirs_core/src/ast/string.rs @@ -48,4 +48,12 @@ impl<'c> Op<'c> for StringOp<'c> { .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect()) } } + + fn get_annotations(&self) -> Vec { + if let StringOp::Annotated(inner, anno) = self { + inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + } else { + vec![] + } + } } diff --git a/crates/clarirs_py/src/ast/bool.rs b/crates/clarirs_py/src/ast/bool.rs index 9e2c4d2..c82cee2 100644 --- a/crates/clarirs_py/src/ast/bool.rs +++ b/crates/clarirs_py/src/ast/bool.rs @@ -6,6 +6,7 @@ use std::sync::LazyLock; use ast::args::ExtractPyArgs; use dashmap::DashMap; +use pyo3::types::PyBytes; use pyo3::types::PyFrozenSet; use pyo3::types::PyWeakrefMethods; use pyo3::types::PyWeakrefReference; @@ -84,6 +85,17 @@ impl Bool { self.inner.symbolic() } + #[getter] + fn annotations(&self, py: Python) -> PyResult> { + let pickle_loads = py.import_bound("pickle")?.getattr("loads")?; + self.inner + .get_annotations() + .iter() + .map(|a| pickle_loads.call1((PyBytes::new_bound(py, a.value()),))) + .map(|a| a.map(|a| a.unbind())) + .collect() + } + fn hash(&self) -> u64 { self.inner.hash() } @@ -105,6 +117,23 @@ impl Bool { self.inner.is_false() } + fn annotate(&self, py: Python, annotation: Bound) -> Result, ClaripyError> { + let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?; + let annotation_bytes = pickle_dumps + .call1((&annotation,))? + .downcast::()? + .extract::>()?; + let eliminatable = annotation.getattr("eliminatable")?.extract::()?; + let relocatable = annotation.getattr("relocatable")?.extract::()?; + Bool::new( + py, + &GLOBAL_CONTEXT.annotated( + &self.inner, + Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable), + )?, + ) + } + fn __invert__(&self, py: Python) -> Result, ClaripyError> { Bool::new(py, &GLOBAL_CONTEXT.not(&self.inner)?) } diff --git a/crates/clarirs_py/src/ast/bv.rs b/crates/clarirs_py/src/ast/bv.rs index 8c1877f..e708bee 100644 --- a/crates/clarirs_py/src/ast/bv.rs +++ b/crates/clarirs_py/src/ast/bv.rs @@ -6,7 +6,7 @@ use std::sync::LazyLock; use dashmap::DashMap; use num_bigint::BigUint; use pyo3::exceptions::{PyTypeError, PyValueError}; -use pyo3::types::{PyFrozenSet, PyWeakrefReference}; +use pyo3::types::{PyBytes, PyFrozenSet, PyWeakrefReference}; use crate::ast::{And, Not, Or, Xor}; use crate::prelude::*; @@ -84,6 +84,17 @@ impl BV { self.inner.symbolic() } + #[getter] + fn annotations(&self, py: Python) -> PyResult> { + let pickle_loads = py.import_bound("pickle")?.getattr("loads")?; + self.inner + .get_annotations() + .iter() + .map(|a| pickle_loads.call1((PyBytes::new_bound(py, a.value()),))) + .map(|a| a.map(|a| a.unbind())) + .collect() + } + fn hash(&self) -> u64 { self.inner.hash() } @@ -97,6 +108,23 @@ impl BV { self.inner.depth() == 1 } + fn annotate(&self, py: Python, annotation: Bound) -> Result, ClaripyError> { + let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?; + let annotation_bytes = pickle_dumps + .call1((&annotation,))? + .downcast::()? + .extract::>()?; + let eliminatable = annotation.getattr("eliminatable")?.extract::()?; + let relocatable = annotation.getattr("relocatable")?.extract::()?; + BV::new( + py, + &GLOBAL_CONTEXT.annotated( + &self.inner, + Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable), + )?, + ) + } + fn __add__(&self, py: Python, other: CoerceBV) -> Result, ClaripyError> { BV::new( py, diff --git a/crates/clarirs_py/src/ast/fp.rs b/crates/clarirs_py/src/ast/fp.rs index 5bb5b25..1b4f3b1 100644 --- a/crates/clarirs_py/src/ast/fp.rs +++ b/crates/clarirs_py/src/ast/fp.rs @@ -6,7 +6,7 @@ use std::sync::{ }; use dashmap::DashMap; -use pyo3::types::{PyFrozenSet, PyWeakrefReference}; +use pyo3::types::{PyBytes, PyFrozenSet, PyWeakrefReference}; use crate::prelude::*; @@ -167,6 +167,17 @@ impl FP { self.inner.symbolic() } + #[getter] + fn annotations(&self, py: Python) -> PyResult> { + let pickle_loads = py.import_bound("pickle")?.getattr("loads")?; + self.inner + .get_annotations() + .iter() + .map(|a| pickle_loads.call1((PyBytes::new_bound(py, a.value()),))) + .map(|a| a.map(|a| a.unbind())) + .collect() + } + fn hash(&self) -> u64 { self.inner.hash() } @@ -179,6 +190,23 @@ impl FP { fn is_leaf(&self) -> bool { self.inner.depth() == 1 } + + fn annotate(&self, py: Python, annotation: Bound) -> Result, ClaripyError> { + let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?; + let annotation_bytes = pickle_dumps + .call1((&annotation,))? + .downcast::()? + .extract::>()?; + let eliminatable = annotation.getattr("eliminatable")?.extract::()?; + let relocatable = annotation.getattr("relocatable")?.extract::()?; + FP::new( + py, + &GLOBAL_CONTEXT.annotated( + &self.inner, + Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable), + )?, + ) + } } #[pyfunction(signature = (name, sort, explicit_name = false))] diff --git a/crates/clarirs_py/src/ast/string.rs b/crates/clarirs_py/src/ast/string.rs index a56f6f7..a1c9a13 100644 --- a/crates/clarirs_py/src/ast/string.rs +++ b/crates/clarirs_py/src/ast/string.rs @@ -6,7 +6,7 @@ use std::sync::{ }; use dashmap::DashMap; -use pyo3::types::{PyFrozenSet, PyWeakrefReference}; +use pyo3::types::{PyBytes, PyFrozenSet, PyWeakrefReference}; use crate::prelude::*; @@ -82,6 +82,17 @@ impl PyAstString { self.inner.symbolic() } + #[getter] + fn annotations(&self, py: Python) -> PyResult> { + let pickle_loads = py.import_bound("pickle")?.getattr("loads")?; + self.inner + .get_annotations() + .iter() + .map(|a| pickle_loads.call1((PyBytes::new_bound(py, a.value()),))) + .map(|a| a.map(|a| a.unbind())) + .collect() + } + fn hash(&self) -> u64 { self.inner.hash() } @@ -94,6 +105,23 @@ impl PyAstString { fn is_leaf(&self) -> bool { self.inner.depth() == 1 } + + fn annotate(&self, py: Python, annotation: Bound) -> Result, ClaripyError> { + let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?; + let annotation_bytes = pickle_dumps + .call1((&annotation,))? + .downcast::()? + .extract::>()?; + let eliminatable = annotation.getattr("eliminatable")?.extract::()?; + let relocatable = annotation.getattr("relocatable")?.extract::()?; + PyAstString::new( + py, + &GLOBAL_CONTEXT.annotated( + &self.inner, + Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable), + )?, + ) + } } #[pyfunction(signature = (name, explicit_name = false))]