_pylong.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285
  1. """Python implementations of some algorithms for use by longobject.c.
  2. The goal is to provide asymptotically faster algorithms that can be
  3. used for operations on integers with many digits. In those cases, the
  4. performance overhead of the Python implementation is not significant
  5. since the asymptotic behavior is what dominates runtime. Functions
  6. provided by this module should be considered private and not part of any
  7. public API.
  8. Note: for ease of maintainability, please prefer clear code and avoid
  9. "micro-optimizations". This module will only be imported and used for
  10. integers with a huge number of digits. Saving a few microseconds with
  11. tricky or non-obvious code is not worth it. For people looking for
  12. maximum performance, they should use something like gmpy2."""
  13. import re
  14. import decimal
  15. def int_to_decimal(n):
  16. """Asymptotically fast conversion of an 'int' to Decimal."""
  17. # Function due to Tim Peters. See GH issue #90716 for details.
  18. # https://github.com/python/cpython/issues/90716
  19. #
  20. # The implementation in longobject.c of base conversion algorithms
  21. # between power-of-2 and non-power-of-2 bases are quadratic time.
  22. # This function implements a divide-and-conquer algorithm that is
  23. # faster for large numbers. Builds an equal decimal.Decimal in a
  24. # "clever" recursive way. If we want a string representation, we
  25. # apply str to _that_.
  26. D = decimal.Decimal
  27. D2 = D(2)
  28. BITLIM = 128
  29. mem = {}
  30. def w2pow(w):
  31. """Return D(2)**w and store the result. Also possibly save some
  32. intermediate results. In context, these are likely to be reused
  33. across various levels of the conversion to Decimal."""
  34. if (result := mem.get(w)) is None:
  35. if w <= BITLIM:
  36. result = D2**w
  37. elif w - 1 in mem:
  38. result = (t := mem[w - 1]) + t
  39. else:
  40. w2 = w >> 1
  41. # If w happens to be odd, w-w2 is one larger then w2
  42. # now. Recurse on the smaller first (w2), so that it's
  43. # in the cache and the larger (w-w2) can be handled by
  44. # the cheaper `w-1 in mem` branch instead.
  45. result = w2pow(w2) * w2pow(w - w2)
  46. mem[w] = result
  47. return result
  48. def inner(n, w):
  49. if w <= BITLIM:
  50. return D(n)
  51. w2 = w >> 1
  52. hi = n >> w2
  53. lo = n - (hi << w2)
  54. return inner(lo, w2) + inner(hi, w - w2) * w2pow(w2)
  55. with decimal.localcontext() as ctx:
  56. ctx.prec = decimal.MAX_PREC
  57. ctx.Emax = decimal.MAX_EMAX
  58. ctx.Emin = decimal.MIN_EMIN
  59. ctx.traps[decimal.Inexact] = 1
  60. if n < 0:
  61. negate = True
  62. n = -n
  63. else:
  64. negate = False
  65. result = inner(n, n.bit_length())
  66. if negate:
  67. result = -result
  68. return result
  69. def int_to_decimal_string(n):
  70. """Asymptotically fast conversion of an 'int' to a decimal string."""
  71. return str(int_to_decimal(n))
  72. def _str_to_int_inner(s):
  73. """Asymptotically fast conversion of a 'str' to an 'int'."""
  74. # Function due to Bjorn Martinsson. See GH issue #90716 for details.
  75. # https://github.com/python/cpython/issues/90716
  76. #
  77. # The implementation in longobject.c of base conversion algorithms
  78. # between power-of-2 and non-power-of-2 bases are quadratic time.
  79. # This function implements a divide-and-conquer algorithm making use
  80. # of Python's built in big int multiplication. Since Python uses the
  81. # Karatsuba algorithm for multiplication, the time complexity
  82. # of this function is O(len(s)**1.58).
  83. DIGLIM = 2048
  84. mem = {}
  85. def w5pow(w):
  86. """Return 5**w and store the result.
  87. Also possibly save some intermediate results. In context, these
  88. are likely to be reused across various levels of the conversion
  89. to 'int'.
  90. """
  91. if (result := mem.get(w)) is None:
  92. if w <= DIGLIM:
  93. result = 5**w
  94. elif w - 1 in mem:
  95. result = mem[w - 1] * 5
  96. else:
  97. w2 = w >> 1
  98. # If w happens to be odd, w-w2 is one larger then w2
  99. # now. Recurse on the smaller first (w2), so that it's
  100. # in the cache and the larger (w-w2) can be handled by
  101. # the cheaper `w-1 in mem` branch instead.
  102. result = w5pow(w2) * w5pow(w - w2)
  103. mem[w] = result
  104. return result
  105. def inner(a, b):
  106. if b - a <= DIGLIM:
  107. return int(s[a:b])
  108. mid = (a + b + 1) >> 1
  109. return inner(mid, b) + ((inner(a, mid) * w5pow(b - mid)) << (b - mid))
  110. return inner(0, len(s))
  111. def int_from_string(s):
  112. """Asymptotically fast version of PyLong_FromString(), conversion
  113. of a string of decimal digits into an 'int'."""
  114. # PyLong_FromString() has already removed leading +/-, checked for invalid
  115. # use of underscore characters, checked that string consists of only digits
  116. # and underscores, and stripped leading whitespace. The input can still
  117. # contain underscores and have trailing whitespace.
  118. s = s.rstrip().replace('_', '')
  119. return _str_to_int_inner(s)
  120. def str_to_int(s):
  121. """Asymptotically fast version of decimal string to 'int' conversion."""
  122. # FIXME: this doesn't support the full syntax that int() supports.
  123. m = re.match(r'\s*([+-]?)([0-9_]+)\s*', s)
  124. if not m:
  125. raise ValueError('invalid literal for int() with base 10')
  126. v = int_from_string(m.group(2))
  127. if m.group(1) == '-':
  128. v = -v
  129. return v
  130. # Fast integer division, based on code from Mark Dickinson, fast_div.py
  131. # GH-47701. Additional refinements and optimizations by Bjorn Martinsson. The
  132. # algorithm is due to Burnikel and Ziegler, in their paper "Fast Recursive
  133. # Division".
  134. _DIV_LIMIT = 4000
  135. def _div2n1n(a, b, n):
  136. """Divide a 2n-bit nonnegative integer a by an n-bit positive integer
  137. b, using a recursive divide-and-conquer algorithm.
  138. Inputs:
  139. n is a positive integer
  140. b is a positive integer with exactly n bits
  141. a is a nonnegative integer such that a < 2**n * b
  142. Output:
  143. (q, r) such that a = b*q+r and 0 <= r < b.
  144. """
  145. if a.bit_length() - n <= _DIV_LIMIT:
  146. return divmod(a, b)
  147. pad = n & 1
  148. if pad:
  149. a <<= 1
  150. b <<= 1
  151. n += 1
  152. half_n = n >> 1
  153. mask = (1 << half_n) - 1
  154. b1, b2 = b >> half_n, b & mask
  155. q1, r = _div3n2n(a >> n, (a >> half_n) & mask, b, b1, b2, half_n)
  156. q2, r = _div3n2n(r, a & mask, b, b1, b2, half_n)
  157. if pad:
  158. r >>= 1
  159. return q1 << half_n | q2, r
  160. def _div3n2n(a12, a3, b, b1, b2, n):
  161. """Helper function for _div2n1n; not intended to be called directly."""
  162. if a12 >> n == b1:
  163. q, r = (1 << n) - 1, a12 - (b1 << n) + b1
  164. else:
  165. q, r = _div2n1n(a12, b1, n)
  166. r = (r << n | a3) - q * b2
  167. while r < 0:
  168. q -= 1
  169. r += b
  170. return q, r
  171. def _int2digits(a, n):
  172. """Decompose non-negative int a into base 2**n
  173. Input:
  174. a is a non-negative integer
  175. Output:
  176. List of the digits of a in base 2**n in little-endian order,
  177. meaning the most significant digit is last. The most
  178. significant digit is guaranteed to be non-zero.
  179. If a is 0 then the output is an empty list.
  180. """
  181. a_digits = [0] * ((a.bit_length() + n - 1) // n)
  182. def inner(x, L, R):
  183. if L + 1 == R:
  184. a_digits[L] = x
  185. return
  186. mid = (L + R) >> 1
  187. shift = (mid - L) * n
  188. upper = x >> shift
  189. lower = x ^ (upper << shift)
  190. inner(lower, L, mid)
  191. inner(upper, mid, R)
  192. if a:
  193. inner(a, 0, len(a_digits))
  194. return a_digits
  195. def _digits2int(digits, n):
  196. """Combine base-2**n digits into an int. This function is the
  197. inverse of `_int2digits`. For more details, see _int2digits.
  198. """
  199. def inner(L, R):
  200. if L + 1 == R:
  201. return digits[L]
  202. mid = (L + R) >> 1
  203. shift = (mid - L) * n
  204. return (inner(mid, R) << shift) + inner(L, mid)
  205. return inner(0, len(digits)) if digits else 0
  206. def _divmod_pos(a, b):
  207. """Divide a non-negative integer a by a positive integer b, giving
  208. quotient and remainder."""
  209. # Use grade-school algorithm in base 2**n, n = nbits(b)
  210. n = b.bit_length()
  211. a_digits = _int2digits(a, n)
  212. r = 0
  213. q_digits = []
  214. for a_digit in reversed(a_digits):
  215. q_digit, r = _div2n1n((r << n) + a_digit, b, n)
  216. q_digits.append(q_digit)
  217. q_digits.reverse()
  218. q = _digits2int(q_digits, n)
  219. return q, r
  220. def int_divmod(a, b):
  221. """Asymptotically fast replacement for divmod, for 'int'.
  222. Its time complexity is O(n**1.58), where n = #bits(a) + #bits(b).
  223. """
  224. if b == 0:
  225. raise ZeroDivisionError
  226. elif b < 0:
  227. q, r = int_divmod(-a, -b)
  228. return q, -r
  229. elif a < 0:
  230. q, r = int_divmod(~a, b)
  231. return ~q, b + ~r
  232. else:
  233. return _divmod_pos(a, b)