From 6c8ccbdbfbbb2bdade653e0b1c8648caf71ef01b Mon Sep 17 00:00:00 2001 From: Kevin Phoenix Date: Thu, 31 Oct 2024 21:01:33 -0700 Subject: [PATCH 1/2] Add basic python annotation support --- crates/clarirs_core/src/ast/bitvec.rs | 8 +++++++ crates/clarirs_core/src/ast/bool.rs | 8 +++++++ crates/clarirs_core/src/ast/float.rs | 8 +++++++ crates/clarirs_core/src/ast/node.rs | 13 ++++++++++++ crates/clarirs_core/src/ast/op.rs | 2 ++ crates/clarirs_core/src/ast/string.rs | 8 +++++++ crates/clarirs_py/src/ast/bool.rs | 29 ++++++++++++++++++++++++++ crates/clarirs_py/src/ast/bv.rs | 30 ++++++++++++++++++++++++++- crates/clarirs_py/src/ast/fp.rs | 30 ++++++++++++++++++++++++++- crates/clarirs_py/src/ast/string.rs | 30 ++++++++++++++++++++++++++- 10 files changed, 163 insertions(+), 3 deletions(-) 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))] From f32a4d50b029c22a13c906400b294e48826a3db2 Mon Sep 17 00:00:00 2001 From: Kevin Phoenix Date: Thu, 31 Oct 2024 21:08:35 -0700 Subject: [PATCH 2/2] cargo fmt --- crates/clarirs_core/src/ast/bitvec.rs | 6 +++++- crates/clarirs_core/src/ast/bool.rs | 6 +++++- crates/clarirs_core/src/ast/float.rs | 6 +++++- crates/clarirs_core/src/ast/string.rs | 6 +++++- crates/clarirs_py/src/ast/string.rs | 6 +++++- 5 files changed, 25 insertions(+), 5 deletions(-) diff --git a/crates/clarirs_core/src/ast/bitvec.rs b/crates/clarirs_core/src/ast/bitvec.rs index 82e47ca..95cd1d2 100644 --- a/crates/clarirs_core/src/ast/bitvec.rs +++ b/crates/clarirs_core/src/ast/bitvec.rs @@ -100,7 +100,11 @@ impl<'c> Op<'c> for BitVecOp<'c> { fn get_annotations(&self) -> Vec { if let BitVecOp::Annotated(inner, anno) = self { - inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + 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 d41984f..ea283db 100644 --- a/crates/clarirs_core/src/ast/bool.rs +++ b/crates/clarirs_core/src/ast/bool.rs @@ -115,7 +115,11 @@ impl<'c> Op<'c> for BooleanOp<'c> { fn get_annotations(&self) -> Vec { if let BooleanOp::Annotated(inner, anno) = self { - inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + 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 b6e94ce..ec7fd71 100644 --- a/crates/clarirs_core/src/ast/float.rs +++ b/crates/clarirs_core/src/ast/float.rs @@ -58,7 +58,11 @@ impl<'c> Op<'c> for FloatOp<'c> { fn get_annotations(&self) -> Vec { if let FloatOp::Annotated(inner, anno) = self { - inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + inner + .get_annotations() + .into_iter() + .chain(vec![anno.clone()]) + .collect() } else { vec![] } diff --git a/crates/clarirs_core/src/ast/string.rs b/crates/clarirs_core/src/ast/string.rs index 26a601e..7cff10e 100644 --- a/crates/clarirs_core/src/ast/string.rs +++ b/crates/clarirs_core/src/ast/string.rs @@ -51,7 +51,11 @@ impl<'c> Op<'c> for StringOp<'c> { fn get_annotations(&self) -> Vec { if let StringOp::Annotated(inner, anno) = self { - inner.get_annotations().into_iter().chain(vec![anno.clone()]).collect() + inner + .get_annotations() + .into_iter() + .chain(vec![anno.clone()]) + .collect() } else { vec![] } diff --git a/crates/clarirs_py/src/ast/string.rs b/crates/clarirs_py/src/ast/string.rs index a1c9a13..279c0a1 100644 --- a/crates/clarirs_py/src/ast/string.rs +++ b/crates/clarirs_py/src/ast/string.rs @@ -106,7 +106,11 @@ impl PyAstString { self.inner.depth() == 1 } - fn annotate(&self, py: Python, annotation: Bound) -> Result, ClaripyError> { + 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,))?