Add basic python annotation support

crates/clarirs_core/src/ast/bitvec.rs

@@ -97,4 +97,16 @@ impl<'c> Op<'c> for BitVecOp<'c> {
                 .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect())
         }
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        if let BitVecOp::Annotated(inner, anno) = self {
+            inner
+                .get_annotations()
+                .into_iter()
+                .chain(vec![anno.clone()])
+                .collect()
+        } else {
+            vec![]
+        }
+    }
 }

crates/clarirs_core/src/ast/bool.rs

@@ -112,4 +112,16 @@ impl<'c> Op<'c> for BooleanOp<'c> {
                 .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect())
         }
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        if let BooleanOp::Annotated(inner, anno) = self {
+            inner
+                .get_annotations()
+                .into_iter()
+                .chain(vec![anno.clone()])
+                .collect()
+        } else {
+            vec![]
+        }
+    }
 }

crates/clarirs_core/src/ast/float.rs

@@ -55,4 +55,16 @@ impl<'c> Op<'c> for FloatOp<'c> {
                 .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect())
         }
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        if let FloatOp::Annotated(inner, anno) = self {
+            inner
+                .get_annotations()
+                .into_iter()
+                .chain(vec![anno.clone()])
+                .collect()
+        } else {
+            vec![]
+        }
+    }
 }

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<String> {
         self.variables.clone()
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        self.op().get_annotations()
+    }
 }
 
 pub type AstRef<'c, Op> = Arc<AstNode<'c, Op>>;
@@ -161,6 +165,15 @@ impl<'c> Op<'c> for VarAst<'c> {
         }
         .clone()
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        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> {

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<Annotation>;
 }

crates/clarirs_core/src/ast/string.rs

@@ -48,4 +48,16 @@ impl<'c> Op<'c> for StringOp<'c> {
                 .fold(HashSet::new(), |acc, x| acc.union(&x).cloned().collect())
         }
     }
+
+    fn get_annotations(&self) -> Vec<Annotation> {
+        if let StringOp::Annotated(inner, anno) = self {
+            inner
+                .get_annotations()
+                .into_iter()
+                .chain(vec![anno.clone()])
+                .collect()
+        } else {
+            vec![]
+        }
+    }
 }

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<Vec<PyObject>> {
+        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<PyAny>) -> Result<Py<Bool>, ClaripyError> {
+        let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
+        let annotation_bytes = pickle_dumps
+            .call1((&annotation,))?
+            .downcast::<PyBytes>()?
+            .extract::<Vec<u8>>()?;
+        let eliminatable = annotation.getattr("eliminatable")?.extract::<bool>()?;
+        let relocatable = annotation.getattr("relocatable")?.extract::<bool>()?;
+        Bool::new(
+            py,
+            &GLOBAL_CONTEXT.annotated(
+                &self.inner,
+                Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable),
+            )?,
+        )
+    }
+
     fn __invert__(&self, py: Python) -> Result<Py<Bool>, ClaripyError> {
         Bool::new(py, &GLOBAL_CONTEXT.not(&self.inner)?)
     }

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<Vec<PyObject>> {
+        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<PyAny>) -> Result<Py<BV>, ClaripyError> {
+        let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
+        let annotation_bytes = pickle_dumps
+            .call1((&annotation,))?
+            .downcast::<PyBytes>()?
+            .extract::<Vec<u8>>()?;
+        let eliminatable = annotation.getattr("eliminatable")?.extract::<bool>()?;
+        let relocatable = annotation.getattr("relocatable")?.extract::<bool>()?;
+        BV::new(
+            py,
+            &GLOBAL_CONTEXT.annotated(
+                &self.inner,
+                Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable),
+            )?,
+        )
+    }
+
     fn __add__(&self, py: Python, other: CoerceBV) -> Result<Py<BV>, ClaripyError> {
         BV::new(
             py,

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<Vec<PyObject>> {
+        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<PyAny>) -> Result<Py<FP>, ClaripyError> {
+        let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
+        let annotation_bytes = pickle_dumps
+            .call1((&annotation,))?
+            .downcast::<PyBytes>()?
+            .extract::<Vec<u8>>()?;
+        let eliminatable = annotation.getattr("eliminatable")?.extract::<bool>()?;
+        let relocatable = annotation.getattr("relocatable")?.extract::<bool>()?;
+        FP::new(
+            py,
+            &GLOBAL_CONTEXT.annotated(
+                &self.inner,
+                Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable),
+            )?,
+        )
+    }
 }
 
 #[pyfunction(signature = (name, sort, explicit_name = false))]

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<Vec<PyObject>> {
+        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,27 @@ impl PyAstString {
     fn is_leaf(&self) -> bool {
         self.inner.depth() == 1
     }
+
+    fn annotate(
+        &self,
+        py: Python,
+        annotation: Bound<PyAny>,
+    ) -> Result<Py<PyAstString>, ClaripyError> {
+        let pickle_dumps = py.import_bound("pickle")?.getattr("dumps")?;
+        let annotation_bytes = pickle_dumps
+            .call1((&annotation,))?
+            .downcast::<PyBytes>()?
+            .extract::<Vec<u8>>()?;
+        let eliminatable = annotation.getattr("eliminatable")?.extract::<bool>()?;
+        let relocatable = annotation.getattr("relocatable")?.extract::<bool>()?;
+        PyAstString::new(
+            py,
+            &GLOBAL_CONTEXT.annotated(
+                &self.inner,
+                Annotation::new("".to_string(), annotation_bytes, eliminatable, relocatable),
+            )?,
+        )
+    }
 }
 
 #[pyfunction(signature = (name, explicit_name = false))]