Skip to content

Commit 5c8ff18

Browse files
twizmwazintwizmwazin
committed
Move Base code to src/ast/base.rs
1 parent 4f1c6e0 commit 5c8ff18

File tree

3 files changed

+283
-279
lines changed

3 files changed

+283
-279
lines changed

Diff for: ‎src/ast/base.rs

+279
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
use std::borrow::Cow;
2+
3+
use pyo3::{
4+
prelude::*,
5+
types::{PyAnyMethods, PyBool, PyBytes, PyDict, PyFloat, PyInt, PyString, PyTuple},
6+
};
7+
8+
#[pyclass(subclass, weakref)]
9+
pub struct Base {
10+
// Hashcons
11+
#[pyo3(get, set)]
12+
op: String,
13+
#[pyo3(get, set)]
14+
args: Py<PyTuple>,
15+
#[pyo3(get, set)]
16+
length: PyObject,
17+
#[pyo3(get, set)]
18+
variables: PyObject, // TODO: This should be a HashSet, leave opaque for now
19+
#[pyo3(get, set)]
20+
symbolic: bool,
21+
#[pyo3(get, set)]
22+
annotations: Py<PyTuple>,
23+
24+
// Not Hashcons
25+
#[pyo3(get, set)]
26+
simplifiable: Option<PyObject>,
27+
#[pyo3(get, set)]
28+
depth: Option<PyObject>,
29+
30+
#[pyo3(get, set)]
31+
_hash: Option<isize>,
32+
#[pyo3(get, set)]
33+
_simplified: Option<PyObject>,
34+
#[pyo3(get, set)]
35+
_cache_key: Option<PyObject>,
36+
#[pyo3(get, set)]
37+
_cached_encoded_name: Option<PyObject>,
38+
#[pyo3(get, set)]
39+
_errored: Option<PyObject>,
40+
#[pyo3(get, set)]
41+
_eager_backends: Option<PyObject>,
42+
#[pyo3(get, set)]
43+
_excavated: Option<PyObject>,
44+
#[pyo3(get, set)]
45+
_burrowed: Option<PyObject>,
46+
#[pyo3(get, set)]
47+
_uninitialized: Option<PyObject>,
48+
#[pyo3(get, set)]
49+
_uc_alloc_depth: Option<PyObject>,
50+
#[pyo3(get, set)]
51+
_uneliminatable_annotations: Option<PyObject>,
52+
#[pyo3(get, set)]
53+
_relocatable_annotations: Option<PyObject>,
54+
}
55+
56+
#[pymethods]
57+
impl Base {
58+
#[new]
59+
#[pyo3(signature = (op, args, length, variables, symbolic, annotations))]
60+
fn new(
61+
op: String,
62+
args: Py<PyTuple>,
63+
length: PyObject,
64+
variables: PyObject,
65+
symbolic: bool,
66+
annotations: Py<PyTuple>,
67+
) -> PyResult<Self> {
68+
Ok(Base {
69+
op,
70+
args,
71+
length,
72+
variables,
73+
symbolic,
74+
annotations,
75+
76+
simplifiable: None,
77+
depth: None,
78+
79+
_hash: None,
80+
_simplified: None,
81+
_cache_key: None,
82+
_cached_encoded_name: None,
83+
_errored: None,
84+
_eager_backends: None,
85+
_excavated: None,
86+
_burrowed: None,
87+
_uninitialized: None,
88+
_uc_alloc_depth: None,
89+
_uneliminatable_annotations: None,
90+
_relocatable_annotations: None,
91+
})
92+
}
93+
94+
#[staticmethod]
95+
fn _arg_serialize<'py>(
96+
py: Python<'py>,
97+
arg: &Bound<'_, PyAny>,
98+
) -> PyResult<Option<Cow<'py, [u8]>>> {
99+
if arg.is_none() {
100+
return Ok(Some(Cow::from(vec![b'\x0f'])));
101+
}
102+
if arg.is(&*PyBool::new_bound(py, true)) {
103+
return Ok(Some(Cow::from(vec![b'\x1f'])));
104+
}
105+
if arg.is(&*PyBool::new_bound(py, false)) {
106+
return Ok(Some(Cow::from(vec![b'\x2e'])));
107+
}
108+
if arg.is_instance(&py.get_type_bound::<PyInt>())? {
109+
let arg = arg.downcast::<PyInt>()?.extract::<i128>()?;
110+
let mut result = Vec::new();
111+
if arg < 0 {
112+
result.push(b'-');
113+
if arg >= -0x7FFF {
114+
result.extend_from_slice(&(arg as i16).to_le_bytes());
115+
} else if arg >= -0x7FFF_FFFF {
116+
result.extend_from_slice(&(arg as i32).to_le_bytes());
117+
} else if arg >= -0x7FFF_FFFF_FFFF_FFFF {
118+
result.extend_from_slice(&(arg as i64).to_le_bytes());
119+
} else {
120+
return Ok(None);
121+
}
122+
} else {
123+
if arg <= 0xFFFF {
124+
result.extend_from_slice(&(arg as i16).to_le_bytes());
125+
} else if arg <= 0xFFFF_FFFF {
126+
result.extend_from_slice(&(arg as i32).to_le_bytes());
127+
} else if arg <= 0xFFFF_FFFF_FFFF_FFFF {
128+
result.extend_from_slice(&(arg as i64).to_le_bytes());
129+
} else {
130+
return Ok(None);
131+
}
132+
}
133+
return Ok(Some(Cow::from(result)));
134+
}
135+
if arg.is_instance(&py.get_type_bound::<PyString>())? {
136+
let arg: String = arg.downcast::<PyString>()?.extract()?;
137+
return Ok(Some(Cow::from(arg.into_bytes())));
138+
}
139+
if arg.is_instance(&py.get_type_bound::<PyFloat>())? {
140+
return Ok(Some(Cow::from(Vec::from(
141+
arg.downcast::<PyFloat>()?.extract::<f32>()?.to_le_bytes(),
142+
))));
143+
}
144+
if arg.is_instance(&py.get_type_bound::<PyTuple>())? {
145+
let mut result = Vec::new();
146+
for item in arg.downcast::<PyTuple>()?.iter() {
147+
if let Some(sub_result) = Self::_arg_serialize(py, &item)? {
148+
result.extend(sub_result.iter());
149+
} else {
150+
return Ok(None); // Do we really want to return None here?
151+
}
152+
}
153+
return Ok(Some(Cow::from(result)));
154+
}
155+
Ok(None)
156+
}
157+
158+
#[staticmethod]
159+
fn _ast_serialize<'py>(
160+
py: Python<'py>,
161+
op: String,
162+
args_tuple: &Bound<'_, PyTuple>,
163+
keywords: &Bound<'_, PyDict>, // TODO: This should be a struct or seperate args
164+
) -> PyResult<Option<Cow<'py, [u8]>>> {
165+
let serailized_args = match Base::_arg_serialize(py, args_tuple)? {
166+
Some(args) => args,
167+
None => return Ok(None),
168+
};
169+
170+
let length = match keywords.contains("length")? {
171+
true => match Base::_arg_serialize(py, &keywords.get_item("length")?.unwrap())? {
172+
Some(length) => length,
173+
None => return Ok(None),
174+
},
175+
false => Cow::from(Vec::from(b"none")),
176+
};
177+
178+
// get_item was unchecked in the python version too
179+
let variables = (keywords.get_item("variables")?.unwrap().hash()? as u64).to_le_bytes();
180+
// this one was unchecked too
181+
let symbolic = match keywords.get_item("symbolic")?.unwrap().is_truthy()? {
182+
true => Cow::from(Vec::from(b"\x01")),
183+
false => Cow::from(Vec::from(b"\x00")),
184+
};
185+
let annotations = match keywords.get_item("annotations")? {
186+
Some(item) => Cow::from(Vec::from((item.hash()? as u64).to_le_bytes())),
187+
None => Cow::from(Vec::from(b"\xf9")),
188+
};
189+
190+
Ok(Some(Cow::from(
191+
[
192+
op.as_bytes(),
193+
&serailized_args,
194+
&length,
195+
&variables,
196+
&symbolic,
197+
&annotations,
198+
]
199+
.concat(),
200+
)))
201+
}
202+
203+
#[staticmethod]
204+
fn _calc_hash<'py>(
205+
py: Python<'py>,
206+
op: String,
207+
args: &Bound<PyTuple>,
208+
keywords: Bound<PyDict>,
209+
) -> PyResult<isize> {
210+
let mut args_tuple = Vec::new();
211+
for arg in args.iter() {
212+
if arg.is_instance(&py.get_type_bound::<PyInt>())?
213+
|| arg.is_instance(&py.get_type_bound::<PyFloat>())?
214+
{
215+
args_tuple.push(arg);
216+
} else {
217+
if arg.hasattr("_hash")? {
218+
args_tuple.push(
219+
arg.getattr("_hash")?
220+
.downcast::<PyInt>()
221+
.unwrap()
222+
.clone()
223+
.into_any(),
224+
);
225+
} else {
226+
args_tuple.push(
227+
// Call hash on the object
228+
arg.call_method0("__hash__")?
229+
.downcast::<PyInt>()
230+
.unwrap()
231+
.clone()
232+
.into_any(),
233+
);
234+
}
235+
}
236+
}
237+
238+
let to_hash = match Base::_ast_serialize(py, op.clone(), &args, &keywords)? {
239+
Some(to_hash) => to_hash,
240+
None => {
241+
let hash_tuple: Bound<PyTuple> = PyTuple::new_bound(
242+
py,
243+
vec![
244+
op.to_object(py).bind(py).as_ref(),
245+
args_tuple.to_object(py).bind(py).as_ref(),
246+
keywords
247+
.get_item("length")?
248+
.unwrap_or(py.None().into_bound(py))
249+
.str()?
250+
.as_ref(),
251+
keywords
252+
.get_item("variables")?
253+
.unwrap() // Unchecked unwrap in python version
254+
.hash()?
255+
.to_object(py)
256+
.bind(py),
257+
keywords.get_item("symbolic")?.unwrap().as_ref(), // Unchecked unwrap in python version
258+
keywords
259+
.get_item("annotations")?
260+
.unwrap_or(py.None().into_bound(py))
261+
.hash()?
262+
.to_object(py)
263+
.bind(py),
264+
],
265+
);
266+
Cow::from(Vec::from(
267+
py.import_bound("pickle")?
268+
.getattr("dumps")?
269+
.call1(PyTuple::new_bound(py, vec![&hash_tuple]))?
270+
.downcast_into::<PyBytes>()?
271+
.as_bytes(),
272+
))
273+
}
274+
};
275+
Ok(isize::from_be_bytes(
276+
(md5::compute(to_hash).0)[..8].try_into().unwrap(),
277+
))
278+
}
279+
}

Diff for: ‎src/ast/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
pub mod base;

0 commit comments

Comments
 (0)