rust.py 21 KB

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