1use ocas_domain::{AlgebraicElement, AlgebraicNumberField, Rational, RationalDomain};
22use ocas_poly::DenseUnivariatePolynomial;
23use pyo3::exceptions::{PyTypeError, PyValueError};
24use pyo3::prelude::*;
25
26#[pyclass(name = "AlgebraicExtension")]
33pub struct PyAlgebraicExtension {
34 pub(crate) field: AlgebraicNumberField,
35}
36
37#[pyclass(name = "AlgebraicElement")]
40pub struct PyAlgebraicElement {
41 pub(crate) elem: AlgebraicElement<Rational>,
42}
43
44#[pyclass(name = "AlgebraicPolynomial", skip_from_py_object)]
46#[derive(Clone)]
47pub struct PyAlgebraicPolynomial {
48 pub(crate) inner: DenseUnivariatePolynomial<AlgebraicNumberField>,
49}
50
51#[pyclass(name = "AlgebraicFactor", skip_from_py_object)]
54pub struct PyAlgebraicFactor {
55 #[pyo3(get)]
56 pub factor: PyAlgebraicPolynomial,
57 #[pyo3(get)]
58 pub multiplicity: usize,
59}
60
61fn py_to_rational(obj: &Bound<'_, PyAny>) -> PyResult<Rational> {
67 if let Ok(n) = obj.extract::<i64>() {
68 Ok(Rational::new(n, 1))
69 } else if let Ok((num, den)) = obj.extract::<(i64, i64)>() {
70 if den == 0 {
71 Err(PyValueError::new_err("rational denominator cannot be zero"))
72 } else {
73 Ok(Rational::new(num, den))
74 }
75 } else {
76 Err(PyTypeError::new_err("expected int or (num, denom) tuple"))
77 }
78}
79
80fn parse_min_poly(coeffs: &Bound<'_, PyAny>) -> PyResult<Vec<Rational>> {
82 let iter = coeffs.try_iter().map_err(|_| {
83 PyTypeError::new_err(
84 "minimal polynomial coefficients must be a list of ints or (num, denom) tuples",
85 )
86 })?;
87 iter.map(|c| py_to_rational(&c?)).collect()
88}
89
90fn py_to_anf_element(
97 field: &AlgebraicNumberField,
98 obj: &Bound<'_, PyAny>,
99) -> PyResult<AlgebraicElement<Rational>> {
100 if let Ok(r) = py_to_rational(obj) {
101 return Ok(field.from_base(r));
102 }
103 if let Ok(elem) = obj.extract::<PyRef<'_, PyAlgebraicElement>>() {
104 return Ok(elem.elem.clone());
105 }
106 if obj.try_iter().is_ok() {
107 let cs: PyResult<Vec<Rational>> = obj
108 .try_iter()
109 .unwrap()
110 .map(|c| py_to_rational(&c?))
111 .collect();
112 return Ok(field.element(cs?));
113 }
114 Err(PyTypeError::new_err(
115 "coefficient must be int, (num, denom), list, or AlgebraicElement",
116 ))
117}
118
119#[pymethods]
124impl PyAlgebraicExtension {
125 #[new]
130 fn new(min_poly: &Bound<'_, PyAny>) -> PyResult<Self> {
131 let coeffs = parse_min_poly(min_poly)?;
132 if coeffs.len() < 2 {
133 return Err(PyValueError::new_err(
134 "minimal polynomial must have degree at least 1",
135 ));
136 }
137 if coeffs.last() != Some(&Rational::new(1, 1)) {
138 return Err(PyValueError::new_err("minimal polynomial must be monic"));
139 }
140 Ok(Self {
141 field: AlgebraicNumberField::new(RationalDomain, coeffs),
142 })
143 }
144
145 fn extension_degree(&self) -> usize {
147 self.field.extension_degree()
148 }
149
150 fn alpha(&self) -> PyAlgebraicElement {
152 PyAlgebraicElement {
153 elem: self.field.alpha(),
154 }
155 }
156
157 #[allow(clippy::wrong_self_convention)]
159 fn from_base(&self, c: &Bound<'_, PyAny>) -> PyResult<PyAlgebraicElement> {
160 let r = py_to_rational(c)?;
161 Ok(PyAlgebraicElement {
162 elem: self.field.from_base(r),
163 })
164 }
165
166 fn element(&self, coeffs: &Bound<'_, PyAny>) -> PyResult<PyAlgebraicElement> {
168 let iter = coeffs.try_iter().map_err(|_| {
169 PyTypeError::new_err(
170 "element coefficients must be a list of ints or (num, denom) tuples",
171 )
172 })?;
173 let cs: PyResult<Vec<Rational>> = iter.map(|c| py_to_rational(&c?)).collect();
174 Ok(PyAlgebraicElement {
175 elem: self.field.element(cs?),
176 })
177 }
178
179 fn __repr__(&self) -> String {
180 format!("AlgebraicExtension(deg={})", self.field.extension_degree())
181 }
182}
183
184#[pymethods]
189impl PyAlgebraicElement {
190 fn coeffs(&self) -> Vec<String> {
193 self.elem.coeffs().iter().map(|c| c.to_string()).collect()
194 }
195
196 fn __str__(&self) -> String {
197 format!("{}", self.elem)
198 }
199
200 fn __repr__(&self) -> String {
201 format!("AlgebraicElement({})", self.elem)
202 }
203}
204
205#[pymethods]
210impl PyAlgebraicPolynomial {
211 #[new]
218 fn new(field: PyRef<'_, PyAlgebraicExtension>, coeffs: &Bound<'_, PyAny>) -> PyResult<Self> {
219 let f = &field.field;
220 let iter = coeffs
221 .try_iter()
222 .map_err(|_| PyTypeError::new_err("polynomial coefficients must be a list"))?;
223 let mut out = Vec::new();
224 for c in iter {
225 out.push(py_to_anf_element(f, &c?)?);
226 }
227 Ok(Self {
228 inner: DenseUnivariatePolynomial::from_coeffs(f.clone(), out),
229 })
230 }
231
232 fn degree(&self) -> Option<usize> {
234 self.inner.degree()
235 }
236
237 fn len(&self) -> usize {
239 self.inner.coeffs().len()
240 }
241
242 fn is_zero(&self) -> bool {
244 self.inner.is_zero()
245 }
246
247 fn coeffs(&self) -> Vec<Vec<String>> {
251 self.inner
252 .coeffs()
253 .iter()
254 .map(|c| c.coeffs().iter().map(|r| r.to_string()).collect())
255 .collect()
256 }
257
258 fn __str__(&self) -> String {
259 format!("{}", PolyDisplay(&self.inner))
260 }
261
262 fn factor(&self) -> Vec<PyAlgebraicFactor> {
264 self.inner
265 .factor()
266 .into_iter()
267 .map(|(f, m)| PyAlgebraicFactor {
268 factor: PyAlgebraicPolynomial { inner: f },
269 multiplicity: m,
270 })
271 .collect()
272 }
273}
274
275struct PolyDisplay<'a>(&'a DenseUnivariatePolynomial<AlgebraicNumberField>);
281
282impl std::fmt::Display for PolyDisplay<'_> {
283 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
284 let coeffs = self.0.coeffs();
285 if coeffs.is_empty() {
286 return write!(f, "0");
287 }
288 let mut first = true;
289 for (i, c) in coeffs.iter().enumerate() {
290 if c.coeffs().is_empty() {
292 continue;
293 }
294 if !first {
295 write!(f, " + ")?;
296 }
297 first = false;
298 match i {
299 0 => write!(f, "({})", c)?,
300 1 => write!(f, "({})*x", c)?,
301 _ => write!(f, "({})*x^{}", c, i)?,
302 }
303 }
304 if first {
305 write!(f, "0")?;
306 }
307 Ok(())
308 }
309}