Coverage for /usr/lib/python3/dist-packages/sympy/printing/rust.py: 25%
187 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2Rust code printer
4The `RustCodePrinter` converts SymPy expressions into Rust expressions.
6A complete code generator, which uses `rust_code` extensively, can be found
7in `sympy.utilities.codegen`. The `codegen` module can be used to generate
8complete source code files.
10"""
12# Possible Improvement
13#
14# * make sure we follow Rust Style Guidelines_
15# * make use of pattern matching
16# * better support for reference
17# * generate generic code and use trait to make sure they have specific methods
18# * use crates_ to get more math support
19# - num_
20# + BigInt_, BigUint_
21# + Complex_
22# + Rational64_, Rational32_, BigRational_
23#
24# .. _crates: https://crates.io/
25# .. _Guidelines: https://github.com/rust-lang/rust/tree/master/src/doc/style
26# .. _num: http://rust-num.github.io/num/num/
27# .. _BigInt: http://rust-num.github.io/num/num/bigint/struct.BigInt.html
28# .. _BigUint: http://rust-num.github.io/num/num/bigint/struct.BigUint.html
29# .. _Complex: http://rust-num.github.io/num/num/complex/struct.Complex.html
30# .. _Rational32: http://rust-num.github.io/num/num/rational/type.Rational32.html
31# .. _Rational64: http://rust-num.github.io/num/num/rational/type.Rational64.html
32# .. _BigRational: http://rust-num.github.io/num/num/rational/type.BigRational.html
34from __future__ import annotations
35from typing import Any
37from sympy.core import S, Rational, Float, Lambda
38from sympy.core.numbers import equal_valued
39from sympy.printing.codeprinter import CodePrinter
41# Rust's methods for integer and float can be found at here :
42#
43# * `Rust - Primitive Type f64 <https://doc.rust-lang.org/std/primitive.f64.html>`_
44# * `Rust - Primitive Type i64 <https://doc.rust-lang.org/std/primitive.i64.html>`_
45#
46# Function Style :
47#
48# 1. args[0].func(args[1:]), method with arguments
49# 2. args[0].func(), method without arguments
50# 3. args[1].func(), method without arguments (e.g. (e, x) => x.exp())
51# 4. func(args), function with arguments
53# dictionary mapping SymPy function to (argument_conditions, Rust_function).
54# Used in RustCodePrinter._print_Function(self)
56# f64 method in Rust
57known_functions = {
58 # "": "is_nan",
59 # "": "is_infinite",
60 # "": "is_finite",
61 # "": "is_normal",
62 # "": "classify",
63 "floor": "floor",
64 "ceiling": "ceil",
65 # "": "round",
66 # "": "trunc",
67 # "": "fract",
68 "Abs": "abs",
69 "sign": "signum",
70 # "": "is_sign_positive",
71 # "": "is_sign_negative",
72 # "": "mul_add",
73 "Pow": [(lambda base, exp: equal_valued(exp, -1), "recip", 2), # 1.0/x
74 (lambda base, exp: equal_valued(exp, 0.5), "sqrt", 2), # x ** 0.5
75 (lambda base, exp: equal_valued(exp, -0.5), "sqrt().recip", 2), # 1/(x ** 0.5)
76 (lambda base, exp: exp == Rational(1, 3), "cbrt", 2), # x ** (1/3)
77 (lambda base, exp: equal_valued(base, 2), "exp2", 3), # 2 ** x
78 (lambda base, exp: exp.is_integer, "powi", 1), # x ** y, for i32
79 (lambda base, exp: not exp.is_integer, "powf", 1)], # x ** y, for f64
80 "exp": [(lambda exp: True, "exp", 2)], # e ** x
81 "log": "ln",
82 # "": "log", # number.log(base)
83 # "": "log2",
84 # "": "log10",
85 # "": "to_degrees",
86 # "": "to_radians",
87 "Max": "max",
88 "Min": "min",
89 # "": "hypot", # (x**2 + y**2) ** 0.5
90 "sin": "sin",
91 "cos": "cos",
92 "tan": "tan",
93 "asin": "asin",
94 "acos": "acos",
95 "atan": "atan",
96 "atan2": "atan2",
97 # "": "sin_cos",
98 # "": "exp_m1", # e ** x - 1
99 # "": "ln_1p", # ln(1 + x)
100 "sinh": "sinh",
101 "cosh": "cosh",
102 "tanh": "tanh",
103 "asinh": "asinh",
104 "acosh": "acosh",
105 "atanh": "atanh",
106 "sqrt": "sqrt", # To enable automatic rewrites
107}
109# i64 method in Rust
110# known_functions_i64 = {
111# "": "min_value",
112# "": "max_value",
113# "": "from_str_radix",
114# "": "count_ones",
115# "": "count_zeros",
116# "": "leading_zeros",
117# "": "trainling_zeros",
118# "": "rotate_left",
119# "": "rotate_right",
120# "": "swap_bytes",
121# "": "from_be",
122# "": "from_le",
123# "": "to_be", # to big endian
124# "": "to_le", # to little endian
125# "": "checked_add",
126# "": "checked_sub",
127# "": "checked_mul",
128# "": "checked_div",
129# "": "checked_rem",
130# "": "checked_neg",
131# "": "checked_shl",
132# "": "checked_shr",
133# "": "checked_abs",
134# "": "saturating_add",
135# "": "saturating_sub",
136# "": "saturating_mul",
137# "": "wrapping_add",
138# "": "wrapping_sub",
139# "": "wrapping_mul",
140# "": "wrapping_div",
141# "": "wrapping_rem",
142# "": "wrapping_neg",
143# "": "wrapping_shl",
144# "": "wrapping_shr",
145# "": "wrapping_abs",
146# "": "overflowing_add",
147# "": "overflowing_sub",
148# "": "overflowing_mul",
149# "": "overflowing_div",
150# "": "overflowing_rem",
151# "": "overflowing_neg",
152# "": "overflowing_shl",
153# "": "overflowing_shr",
154# "": "overflowing_abs",
155# "Pow": "pow",
156# "Abs": "abs",
157# "sign": "signum",
158# "": "is_positive",
159# "": "is_negnative",
160# }
162# These are the core reserved words in the Rust language. Taken from:
163# http://doc.rust-lang.org/grammar.html#keywords
165reserved_words = ['abstract',
166 'alignof',
167 'as',
168 'become',
169 'box',
170 'break',
171 'const',
172 'continue',
173 'crate',
174 'do',
175 'else',
176 'enum',
177 'extern',
178 'false',
179 'final',
180 'fn',
181 'for',
182 'if',
183 'impl',
184 'in',
185 'let',
186 'loop',
187 'macro',
188 'match',
189 'mod',
190 'move',
191 'mut',
192 'offsetof',
193 'override',
194 'priv',
195 'proc',
196 'pub',
197 'pure',
198 'ref',
199 'return',
200 'Self',
201 'self',
202 'sizeof',
203 'static',
204 'struct',
205 'super',
206 'trait',
207 'true',
208 'type',
209 'typeof',
210 'unsafe',
211 'unsized',
212 'use',
213 'virtual',
214 'where',
215 'while',
216 'yield']
219class RustCodePrinter(CodePrinter):
220 """A printer to convert SymPy expressions to strings of Rust code"""
221 printmethod = "_rust_code"
222 language = "Rust"
224 _default_settings: dict[str, Any] = {
225 'order': None,
226 'full_prec': 'auto',
227 'precision': 17,
228 'user_functions': {},
229 'human': True,
230 'contract': True,
231 'dereference': set(),
232 'error_on_reserved': False,
233 'reserved_word_suffix': '_',
234 'inline': False,
235 }
237 def __init__(self, settings={}):
238 CodePrinter.__init__(self, settings)
239 self.known_functions = dict(known_functions)
240 userfuncs = settings.get('user_functions', {})
241 self.known_functions.update(userfuncs)
242 self._dereference = set(settings.get('dereference', []))
243 self.reserved_words = set(reserved_words)
245 def _rate_index_position(self, p):
246 return p*5
248 def _get_statement(self, codestring):
249 return "%s;" % codestring
251 def _get_comment(self, text):
252 return "// %s" % text
254 def _declare_number_const(self, name, value):
255 return "const %s: f64 = %s;" % (name, value)
257 def _format_code(self, lines):
258 return self.indent_code(lines)
260 def _traverse_matrix_indices(self, mat):
261 rows, cols = mat.shape
262 return ((i, j) for i in range(rows) for j in range(cols))
264 def _get_loop_opening_ending(self, indices):
265 open_lines = []
266 close_lines = []
267 loopstart = "for %(var)s in %(start)s..%(end)s {"
268 for i in indices:
269 # Rust arrays start at 0 and end at dimension-1
270 open_lines.append(loopstart % {
271 'var': self._print(i),
272 'start': self._print(i.lower),
273 'end': self._print(i.upper + 1)})
274 close_lines.append("}")
275 return open_lines, close_lines
277 def _print_caller_var(self, expr):
278 if len(expr.args) > 1:
279 # for something like `sin(x + y + z)`,
280 # make sure we can get '(x + y + z).sin()'
281 # instead of 'x + y + z.sin()'
282 return '(' + self._print(expr) + ')'
283 elif expr.is_number:
284 return self._print(expr, _type=True)
285 else:
286 return self._print(expr)
288 def _print_Function(self, expr):
289 """
290 basic function for printing `Function`
292 Function Style :
294 1. args[0].func(args[1:]), method with arguments
295 2. args[0].func(), method without arguments
296 3. args[1].func(), method without arguments (e.g. (e, x) => x.exp())
297 4. func(args), function with arguments
298 """
300 if expr.func.__name__ in self.known_functions:
301 cond_func = self.known_functions[expr.func.__name__]
302 func = None
303 style = 1
304 if isinstance(cond_func, str):
305 func = cond_func
306 else:
307 for cond, func, style in cond_func:
308 if cond(*expr.args):
309 break
310 if func is not None:
311 if style == 1:
312 ret = "%(var)s.%(method)s(%(args)s)" % {
313 'var': self._print_caller_var(expr.args[0]),
314 'method': func,
315 'args': self.stringify(expr.args[1:], ", ") if len(expr.args) > 1 else ''
316 }
317 elif style == 2:
318 ret = "%(var)s.%(method)s()" % {
319 'var': self._print_caller_var(expr.args[0]),
320 'method': func,
321 }
322 elif style == 3:
323 ret = "%(var)s.%(method)s()" % {
324 'var': self._print_caller_var(expr.args[1]),
325 'method': func,
326 }
327 else:
328 ret = "%(func)s(%(args)s)" % {
329 'func': func,
330 'args': self.stringify(expr.args, ", "),
331 }
332 return ret
333 elif hasattr(expr, '_imp_') and isinstance(expr._imp_, Lambda):
334 # inlined function
335 return self._print(expr._imp_(*expr.args))
336 elif expr.func.__name__ in self._rewriteable_functions:
337 # Simple rewrite to supported function possible
338 target_f, required_fs = self._rewriteable_functions[expr.func.__name__]
339 if self._can_print(target_f) and all(self._can_print(f) for f in required_fs):
340 return self._print(expr.rewrite(target_f))
341 else:
342 return self._print_not_supported(expr)
344 def _print_Pow(self, expr):
345 if expr.base.is_integer and not expr.exp.is_integer:
346 expr = type(expr)(Float(expr.base), expr.exp)
347 return self._print(expr)
348 return self._print_Function(expr)
350 def _print_Float(self, expr, _type=False):
351 ret = super()._print_Float(expr)
352 if _type:
353 return ret + '_f64'
354 else:
355 return ret
357 def _print_Integer(self, expr, _type=False):
358 ret = super()._print_Integer(expr)
359 if _type:
360 return ret + '_i32'
361 else:
362 return ret
364 def _print_Rational(self, expr):
365 p, q = int(expr.p), int(expr.q)
366 return '%d_f64/%d.0' % (p, q)
368 def _print_Relational(self, expr):
369 lhs_code = self._print(expr.lhs)
370 rhs_code = self._print(expr.rhs)
371 op = expr.rel_op
372 return "{} {} {}".format(lhs_code, op, rhs_code)
374 def _print_Indexed(self, expr):
375 # calculate index for 1d array
376 dims = expr.shape
377 elem = S.Zero
378 offset = S.One
379 for i in reversed(range(expr.rank)):
380 elem += expr.indices[i]*offset
381 offset *= dims[i]
382 return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
384 def _print_Idx(self, expr):
385 return expr.label.name
387 def _print_Dummy(self, expr):
388 return expr.name
390 def _print_Exp1(self, expr, _type=False):
391 return "E"
393 def _print_Pi(self, expr, _type=False):
394 return 'PI'
396 def _print_Infinity(self, expr, _type=False):
397 return 'INFINITY'
399 def _print_NegativeInfinity(self, expr, _type=False):
400 return 'NEG_INFINITY'
402 def _print_BooleanTrue(self, expr, _type=False):
403 return "true"
405 def _print_BooleanFalse(self, expr, _type=False):
406 return "false"
408 def _print_bool(self, expr, _type=False):
409 return str(expr).lower()
411 def _print_NaN(self, expr, _type=False):
412 return "NAN"
414 def _print_Piecewise(self, expr):
415 if expr.args[-1].cond != True:
416 # We need the last conditional to be a True, otherwise the resulting
417 # function may not return a result.
418 raise ValueError("All Piecewise expressions must contain an "
419 "(expr, True) statement to be used as a default "
420 "condition. Without one, the generated "
421 "expression may not evaluate to anything under "
422 "some condition.")
423 lines = []
425 for i, (e, c) in enumerate(expr.args):
426 if i == 0:
427 lines.append("if (%s) {" % self._print(c))
428 elif i == len(expr.args) - 1 and c == True:
429 lines[-1] += " else {"
430 else:
431 lines[-1] += " else if (%s) {" % self._print(c)
432 code0 = self._print(e)
433 lines.append(code0)
434 lines.append("}")
436 if self._settings['inline']:
437 return " ".join(lines)
438 else:
439 return "\n".join(lines)
441 def _print_ITE(self, expr):
442 from sympy.functions import Piecewise
443 return self._print(expr.rewrite(Piecewise, deep=False))
445 def _print_MatrixBase(self, A):
446 if A.cols == 1:
447 return "[%s]" % ", ".join(self._print(a) for a in A)
448 else:
449 raise ValueError("Full Matrix Support in Rust need Crates (https://crates.io/keywords/matrix).")
451 def _print_SparseRepMatrix(self, mat):
452 # do not allow sparse matrices to be made dense
453 return self._print_not_supported(mat)
455 def _print_MatrixElement(self, expr):
456 return "%s[%s]" % (expr.parent,
457 expr.j + expr.i*expr.parent.shape[1])
459 def _print_Symbol(self, expr):
461 name = super()._print_Symbol(expr)
463 if expr in self._dereference:
464 return '(*%s)' % name
465 else:
466 return name
468 def _print_Assignment(self, expr):
469 from sympy.tensor.indexed import IndexedBase
470 lhs = expr.lhs
471 rhs = expr.rhs
472 if self._settings["contract"] and (lhs.has(IndexedBase) or
473 rhs.has(IndexedBase)):
474 # Here we check if there is looping to be done, and if so
475 # print the required loops.
476 return self._doprint_loops(rhs, lhs)
477 else:
478 lhs_code = self._print(lhs)
479 rhs_code = self._print(rhs)
480 return self._get_statement("%s = %s" % (lhs_code, rhs_code))
482 def indent_code(self, code):
483 """Accepts a string of code or a list of code lines"""
485 if isinstance(code, str):
486 code_lines = self.indent_code(code.splitlines(True))
487 return ''.join(code_lines)
489 tab = " "
490 inc_token = ('{', '(', '{\n', '(\n')
491 dec_token = ('}', ')')
493 code = [ line.lstrip(' \t') for line in code ]
495 increase = [ int(any(map(line.endswith, inc_token))) for line in code ]
496 decrease = [ int(any(map(line.startswith, dec_token)))
497 for line in code ]
499 pretty = []
500 level = 0
501 for n, line in enumerate(code):
502 if line in ('', '\n'):
503 pretty.append(line)
504 continue
505 level -= decrease[n]
506 pretty.append("%s%s" % (tab*level, line))
507 level += increase[n]
508 return pretty
511def rust_code(expr, assign_to=None, **settings):
512 """Converts an expr to a string of Rust code
514 Parameters
515 ==========
517 expr : Expr
518 A SymPy expression to be converted.
519 assign_to : optional
520 When given, the argument is used as the name of the variable to which
521 the expression is assigned. Can be a string, ``Symbol``,
522 ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
523 line-wrapping, or for expressions that generate multi-line statements.
524 precision : integer, optional
525 The precision for numbers such as pi [default=15].
526 user_functions : dict, optional
527 A dictionary where the keys are string representations of either
528 ``FunctionClass`` or ``UndefinedFunction`` instances and the values
529 are their desired C string representations. Alternatively, the
530 dictionary value can be a list of tuples i.e. [(argument_test,
531 cfunction_string)]. See below for examples.
532 dereference : iterable, optional
533 An iterable of symbols that should be dereferenced in the printed code
534 expression. These would be values passed by address to the function.
535 For example, if ``dereference=[a]``, the resulting code would print
536 ``(*a)`` instead of ``a``.
537 human : bool, optional
538 If True, the result is a single string that may contain some constant
539 declarations for the number symbols. If False, the same information is
540 returned in a tuple of (symbols_to_declare, not_supported_functions,
541 code_text). [default=True].
542 contract: bool, optional
543 If True, ``Indexed`` instances are assumed to obey tensor contraction
544 rules and the corresponding nested loops over indices are generated.
545 Setting contract=False will not generate loops, instead the user is
546 responsible to provide values for the indices in the code.
547 [default=True].
549 Examples
550 ========
552 >>> from sympy import rust_code, symbols, Rational, sin, ceiling, Abs, Function
553 >>> x, tau = symbols("x, tau")
554 >>> rust_code((2*tau)**Rational(7, 2))
555 '8*1.4142135623731*tau.powf(7_f64/2.0)'
556 >>> rust_code(sin(x), assign_to="s")
557 's = x.sin();'
559 Simple custom printing can be defined for certain types by passing a
560 dictionary of {"type" : "function"} to the ``user_functions`` kwarg.
561 Alternatively, the dictionary value can be a list of tuples i.e.
562 [(argument_test, cfunction_string)].
564 >>> custom_functions = {
565 ... "ceiling": "CEIL",
566 ... "Abs": [(lambda x: not x.is_integer, "fabs", 4),
567 ... (lambda x: x.is_integer, "ABS", 4)],
568 ... "func": "f"
569 ... }
570 >>> func = Function('func')
571 >>> rust_code(func(Abs(x) + ceiling(x)), user_functions=custom_functions)
572 '(fabs(x) + x.CEIL()).f()'
574 ``Piecewise`` expressions are converted into conditionals. If an
575 ``assign_to`` variable is provided an if statement is created, otherwise
576 the ternary operator is used. Note that if the ``Piecewise`` lacks a
577 default term, represented by ``(expr, True)`` then an error will be thrown.
578 This is to prevent generating an expression that may not evaluate to
579 anything.
581 >>> from sympy import Piecewise
582 >>> expr = Piecewise((x + 1, x > 0), (x, True))
583 >>> print(rust_code(expr, tau))
584 tau = if (x > 0) {
585 x + 1
586 } else {
587 x
588 };
590 Support for loops is provided through ``Indexed`` types. With
591 ``contract=True`` these expressions will be turned into loops, whereas
592 ``contract=False`` will just print the assignment expression that should be
593 looped over:
595 >>> from sympy import Eq, IndexedBase, Idx
596 >>> len_y = 5
597 >>> y = IndexedBase('y', shape=(len_y,))
598 >>> t = IndexedBase('t', shape=(len_y,))
599 >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
600 >>> i = Idx('i', len_y-1)
601 >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
602 >>> rust_code(e.rhs, assign_to=e.lhs, contract=False)
603 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
605 Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
606 must be provided to ``assign_to``. Note that any expression that can be
607 generated normally can also exist inside a Matrix:
609 >>> from sympy import Matrix, MatrixSymbol
610 >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
611 >>> A = MatrixSymbol('A', 3, 1)
612 >>> print(rust_code(mat, A))
613 A = [x.powi(2), if (x > 0) {
614 x + 1
615 } else {
616 x
617 }, x.sin()];
618 """
620 return RustCodePrinter(settings).doprint(expr, assign_to)
623def print_rust_code(expr, **settings):
624 """Prints Rust representation of the given expression."""
625 print(rust_code(expr, **settings))