Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for explicit_name in python constructors #14

Merged
merged 2 commits into from
Nov 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion crates/clarirs_py/src/ast/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
25 changes: 21 additions & 4 deletions crates/clarirs_py/src/ast/bool.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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")]
Expand All @@ -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)
Expand All @@ -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(),
}),
)?;
Expand Down Expand Up @@ -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]
Expand Down
29 changes: 25 additions & 4 deletions crates/clarirs_py/src/ast/bv.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#![allow(non_snake_case)]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::LazyLock;

use dashmap::DashMap;
Expand All @@ -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")]
Expand All @@ -19,6 +21,14 @@ 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)
Expand All @@ -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(),
Expand Down Expand Up @@ -398,9 +408,20 @@ 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)]
Expand Down
33 changes: 28 additions & 5 deletions crates/clarirs_py/src/ast/fp.rs
Original file line number Diff line number Diff line change
@@ -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<DashMap<u64, Py<PyWeakrefReference>>> = LazyLock::new(DashMap::new);

#[pyclass(name = "RM", module = "clarirs.ast.fp", eq)]
Expand Down Expand Up @@ -100,6 +104,14 @@ 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)
Expand All @@ -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(),
Expand Down Expand Up @@ -169,9 +181,20 @@ 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]
Expand Down
32 changes: 27 additions & 5 deletions crates/clarirs_py/src/ast/string.rs
Original file line number Diff line number Diff line change
@@ -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<DashMap<u64, Py<PyWeakrefReference>>> =
LazyLock::new(DashMap::new);

Expand All @@ -17,6 +21,14 @@ 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)
Expand All @@ -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(),
}),
)?;
Expand Down Expand Up @@ -84,9 +96,19 @@ 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]
Expand Down
Loading