diff --git a/crates/clarirs_py/src/ast/base.rs b/crates/clarirs_py/src/ast/base.rs index 2cf8f6b..9e87a2e 100644 --- a/crates/clarirs_py/src/ast/base.rs +++ b/crates/clarirs_py/src/ast/base.rs @@ -6,20 +6,39 @@ use crate::prelude::*; #[derive(Clone)] pub struct Base { errored: Py, + name: Option, + encoded_name: Option>, } impl Base { pub fn new(py: Python) -> Self { - Base { + Self::new_with_name(py, None) + } + + pub fn new_with_name(py: Python, name: Option) -> Self { + let encoded_name = name.as_ref().map(|name| name.as_bytes().to_vec()); + Self { errored: PySet::empty_bound(py) .expect("Failed to create PySet") .unbind(), + name, + encoded_name, } } } #[pymethods] impl Base { + #[getter] + fn name(&self) -> Option<&str> { + self.name.as_deref() + } + + #[getter] + fn _encoded_name(&self) -> Option<&[u8]> { + self.encoded_name.as_deref() + } + #[getter] fn _errored(&self) -> Py { self.errored.clone() diff --git a/crates/clarirs_py/src/ast/bool.rs b/crates/clarirs_py/src/ast/bool.rs index d19d517..9e2c4d2 100644 --- a/crates/clarirs_py/src/ast/bool.rs +++ b/crates/clarirs_py/src/ast/bool.rs @@ -1,5 +1,7 @@ #![allow(non_snake_case)] +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::sync::LazyLock; use ast::args::ExtractPyArgs; @@ -11,6 +13,7 @@ use pyo3::types::PyWeakrefReference; use crate::ast::{And, Not, Or, Xor}; use crate::prelude::*; +static BOOLS_COUNTER: AtomicUsize = AtomicUsize::new(0); static PY_BOOL_CACHE: LazyLock>> = LazyLock::new(DashMap::new); #[pyclass(extends=Base, subclass, frozen, weakref, module="clarirs.ast.bool")] @@ -20,6 +23,14 @@ pub struct Bool { impl Bool { pub fn new(py: Python, inner: &BoolAst<'static>) -> Result, ClaripyError> { + Self::new_with_name(py, inner, None) + } + + pub fn new_with_name( + py: Python, + inner: &BoolAst<'static>, + name: Option, + ) -> Result, ClaripyError> { if let Some(cache_hit) = PY_BOOL_CACHE.get(&inner.hash()).and_then(|cache_hit| { cache_hit .bind(py) @@ -30,7 +41,7 @@ impl Bool { } else { let this = Py::new( py, - PyClassInitializer::from(Base::new(py)).add_subclass(Bool { + PyClassInitializer::from(Base::new_with_name(py, name)).add_subclass(Bool { inner: inner.clone(), }), )?; @@ -120,9 +131,15 @@ impl Bool { } } -#[pyfunction] -pub fn BoolS(py: Python, name: &str) -> Result, ClaripyError> { - Bool::new(py, &GLOBAL_CONTEXT.bools(name)?) +#[pyfunction(signature = (name, explicit_name = false))] +pub fn BoolS(py: Python, name: &str, explicit_name: bool) -> Result, ClaripyError> { + let name: String = if explicit_name { + name.to_string() + } else { + let counter = BOOLS_COUNTER.fetch_add(1, Ordering::Relaxed); + format!("Bool_{}_{}", name, counter) + }; + Bool::new_with_name(py, &GLOBAL_CONTEXT.bools(&name)?, Some(name.clone())) } #[pyfunction] diff --git a/crates/clarirs_py/src/ast/bv.rs b/crates/clarirs_py/src/ast/bv.rs index 391000b..8c1877f 100644 --- a/crates/clarirs_py/src/ast/bv.rs +++ b/crates/clarirs_py/src/ast/bv.rs @@ -1,5 +1,6 @@ #![allow(non_snake_case)] +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::LazyLock; use dashmap::DashMap; @@ -10,6 +11,7 @@ use pyo3::types::{PyFrozenSet, PyWeakrefReference}; use crate::ast::{And, Not, Or, Xor}; use crate::prelude::*; +static BVS_COUNTER: AtomicUsize = AtomicUsize::new(0); static PY_BV_CACHE: LazyLock>> = LazyLock::new(DashMap::new); #[pyclass(extends=Bits, subclass, frozen, weakref, module="clarirs.ast.bv")] @@ -19,6 +21,14 @@ pub struct BV { impl BV { pub fn new(py: Python, inner: &BitVecAst<'static>) -> Result, ClaripyError> { + Self::new_with_name(py, inner, None) + } + + pub fn new_with_name( + py: Python, + inner: &BitVecAst<'static>, + name: Option, + ) -> Result, ClaripyError> { if let Some(cache_hit) = PY_BV_CACHE.get(&inner.hash()).and_then(|cache_hit| { cache_hit .bind(py) @@ -29,7 +39,7 @@ impl BV { } else { let this = Py::new( py, - PyClassInitializer::from(Base::new(py)) + PyClassInitializer::from(Base::new_with_name(py, name)) .add_subclass(Bits::new()) .add_subclass(BV { inner: inner.clone(), @@ -398,9 +408,20 @@ impl BV { } } -#[pyfunction] -pub fn BVS(py: Python, name: String, size: u32) -> Result, ClaripyError> { - BV::new(py, &GLOBAL_CONTEXT.bvs(&name, size)?) +#[pyfunction(signature = (name, size, explicit_name = false))] +pub fn BVS( + py: Python, + name: String, + size: u32, + explicit_name: bool, +) -> Result, ClaripyError> { + let name: String = if explicit_name { + name.to_string() + } else { + let counter = BVS_COUNTER.fetch_add(1, Ordering::Relaxed); + format!("BV{}_{}_{}", size, name, counter) + }; + BV::new_with_name(py, &GLOBAL_CONTEXT.bvs(&name, size)?, Some(name.clone())) } #[allow(non_snake_case)] diff --git a/crates/clarirs_py/src/ast/fp.rs b/crates/clarirs_py/src/ast/fp.rs index 936b8b0..5bb5b25 100644 --- a/crates/clarirs_py/src/ast/fp.rs +++ b/crates/clarirs_py/src/ast/fp.rs @@ -1,12 +1,16 @@ #![allow(non_snake_case)] -use std::sync::LazyLock; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + LazyLock, +}; use dashmap::DashMap; use pyo3::types::{PyFrozenSet, PyWeakrefReference}; use crate::prelude::*; +static FPS_COUNTER: AtomicUsize = AtomicUsize::new(0); static PY_FP_CACHE: LazyLock>> = LazyLock::new(DashMap::new); #[pyclass(name = "RM", module = "clarirs.ast.fp", eq)] @@ -100,6 +104,14 @@ pub struct FP { impl FP { pub fn new(py: Python, inner: &FloatAst<'static>) -> Result, ClaripyError> { + Self::new_with_name(py, inner, None) + } + + pub fn new_with_name( + py: Python, + inner: &FloatAst<'static>, + name: Option, + ) -> Result, ClaripyError> { if let Some(cache_hit) = PY_FP_CACHE.get(&inner.hash()).and_then(|cache_hit| { cache_hit .bind(py) @@ -110,7 +122,7 @@ impl FP { } else { let this = Py::new( py, - PyClassInitializer::from(Base::new(py)) + PyClassInitializer::from(Base::new_with_name(py, name)) .add_subclass(Bits::new()) .add_subclass(FP { inner: inner.clone(), @@ -169,9 +181,20 @@ impl FP { } } -#[pyfunction] -pub fn FPS(py: Python, name: &str, sort: PyFSort) -> Result, ClaripyError> { - FP::new(py, &GLOBAL_CONTEXT.fps(name, sort)?) +#[pyfunction(signature = (name, sort, explicit_name = false))] +pub fn FPS( + py: Python, + name: &str, + sort: PyFSort, + explicit_name: bool, +) -> Result, ClaripyError> { + let name: String = if explicit_name { + name.to_string() + } else { + let counter = FPS_COUNTER.fetch_add(1, Ordering::Relaxed); + format!("FP{}_{}_{}", sort.0.size(), name, counter) + }; + FP::new_with_name(py, &GLOBAL_CONTEXT.fps(&name, sort)?, Some(name)) } #[pyfunction] diff --git a/crates/clarirs_py/src/ast/string.rs b/crates/clarirs_py/src/ast/string.rs index bd751a6..a56f6f7 100644 --- a/crates/clarirs_py/src/ast/string.rs +++ b/crates/clarirs_py/src/ast/string.rs @@ -1,12 +1,16 @@ #![allow(non_snake_case)] -use std::sync::LazyLock; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + LazyLock, +}; use dashmap::DashMap; use pyo3::types::{PyFrozenSet, PyWeakrefReference}; use crate::prelude::*; +static STRINGS_COUNTER: AtomicUsize = AtomicUsize::new(0); static PY_STRING_CACHE: LazyLock>> = LazyLock::new(DashMap::new); @@ -17,6 +21,14 @@ pub struct PyAstString { impl PyAstString { pub fn new(py: Python, inner: &StringAst<'static>) -> Result, ClaripyError> { + Self::new_with_name(py, inner, None) + } + + pub fn new_with_name( + py: Python, + inner: &StringAst<'static>, + name: Option, + ) -> Result, ClaripyError> { if let Some(cache_hit) = PY_STRING_CACHE.get(&inner.hash()).and_then(|cache_hit| { cache_hit .bind(py) @@ -27,7 +39,7 @@ impl PyAstString { } else { let this = Py::new( py, - PyClassInitializer::from(Base::new(py)).add_subclass(PyAstString { + PyClassInitializer::from(Base::new_with_name(py, name)).add_subclass(PyAstString { inner: inner.clone(), }), )?; @@ -84,9 +96,19 @@ impl PyAstString { } } -#[pyfunction] -pub fn StringS(py: Python, name: &str) -> Result, ClaripyError> { - PyAstString::new(py, &GLOBAL_CONTEXT.strings(name)?) +#[pyfunction(signature = (name, explicit_name = false))] +pub fn StringS( + py: Python, + name: &str, + explicit_name: bool, +) -> Result, ClaripyError> { + let name: String = if explicit_name { + name.to_string() + } else { + let counter = STRINGS_COUNTER.fetch_add(1, Ordering::Relaxed); + format!("String_{}_{}", name, counter) + }; + PyAstString::new_with_name(py, &GLOBAL_CONTEXT.strings(&name)?, Some(name)) } #[pyfunction]