Add support for explicit_name in python constructors

twizmwazin · Nov 1, 2024 · a32c474 · a32c474
1 parent 619aafc
commit a32c474
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 19 deletions.
diff --git 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<PySet>,
+    name: Option<String>,
+    encoded_name: Option<Vec<u8>>,
 }
 
 impl Base {
     pub fn new(py: Python) -> Self {
-        Base {
+        Self::new_with_name(py, None)
+    }
+
+    pub fn new_with_name(py: Python, name: Option<String>) -> 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<PySet> {
         self.errored.clone()

diff --git 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<DashMap<u64, Py<PyWeakrefReference>>> = 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<Py<Bool>, ClaripyError> {
+        Self::new_with_name(py, inner, None)
+    }
+
+    pub fn new_with_name(
+        py: Python,
+        inner: &BoolAst<'static>,
+        name: Option<String>,
+    ) -> Result<Py<Bool>, 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<Py<Bool>, 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<Py<Bool>, 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
@@ -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<DashMap<u64, Py<PyWeakrefReference>>> = LazyLock::new(DashMap::new);
 
 #[pyclass(extends=Bits, subclass, frozen, weakref, module="clarirs.ast.bv")]
@@ -19,6 +21,10 @@ pub struct BV {
 
 impl BV {
     pub fn new(py: Python, inner: &BitVecAst<'static>) -> Result<Py<BV>, ClaripyError> {
+        Self::new_with_name(py, inner, None)
+    }
+
+    pub fn new_with_name(py: Python, inner: &BitVecAst<'static>, name: Option<String>) -> Result<Py<BV>, ClaripyError> {
         if let Some(cache_hit) = PY_BV_CACHE.get(&inner.hash()).and_then(|cache_hit| {
             cache_hit
                 .bind(py)
@@ -29,7 +35,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 +404,15 @@ impl BV {
     }
 }
 
-#[pyfunction]
-pub fn BVS(py: Python, name: String, size: u32) -> Result<Py<BV>, 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<Py<BV>, 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
@@ -1,12 +1,13 @@
 #![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<DashMap<u64, Py<PyWeakrefReference>>> = LazyLock::new(DashMap::new);
 
 #[pyclass(name = "RM", module = "clarirs.ast.fp", eq)]
@@ -100,6 +101,10 @@ pub struct FP {
 
 impl FP {
     pub fn new(py: Python, inner: &FloatAst<'static>) -> Result<Py<FP>, ClaripyError> {
+        Self::new_with_name(py, inner, None)
+    }
+
+    pub fn new_with_name(py: Python, inner: &FloatAst<'static>, name: Option<String>) -> Result<Py<FP>, ClaripyError> {
         if let Some(cache_hit) = PY_FP_CACHE.get(&inner.hash()).and_then(|cache_hit| {
             cache_hit
                 .bind(py)
@@ -110,7 +115,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 +174,15 @@ impl FP {
     }
 }
 
-#[pyfunction]
-pub fn FPS(py: Python, name: &str, sort: PyFSort) -> Result<Py<FP>, 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<Py<FP>, 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
@@ -1,12 +1,13 @@
 #![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<DashMap<u64, Py<PyWeakrefReference>>> =
     LazyLock::new(DashMap::new);
 
@@ -17,6 +18,10 @@ pub struct PyAstString {
 
 impl PyAstString {
     pub fn new(py: Python, inner: &StringAst<'static>) -> Result<Py<PyAstString>, ClaripyError> {
+        Self::new_with_name(py, inner, None)
+    }
+
+    pub fn new_with_name(py: Python, inner: &StringAst<'static>, name: Option<String>) -> Result<Py<PyAstString>, ClaripyError> {
         if let Some(cache_hit) = PY_STRING_CACHE.get(&inner.hash()).and_then(|cache_hit| {
             cache_hit
                 .bind(py)
@@ -27,7 +32,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 +89,15 @@ impl PyAstString {
     }
 }
 
-#[pyfunction]
-pub fn StringS(py: Python, name: &str) -> Result<Py<PyAstString>, 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<Py<PyAstString>, 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]