1# This file is part of Hypothesis, which may be found at
2# https://github.com/HypothesisWorks/hypothesis/
3#
4# Copyright the Hypothesis Authors.
5# Individual contributors are listed in AUTHORS.rst and the git log.
6#
7# This Source Code Form is subject to the terms of the Mozilla Public License,
8# v. 2.0. If a copy of the MPL was not distributed with this file, You can
9# obtain one at https://mozilla.org/MPL/2.0/.
10
11import math
12from collections.abc import Callable, Hashable, Iterable, Sequence
13from dataclasses import dataclass
14from typing import (
15 Literal,
16 TypeAlias,
17 TypedDict,
18 TypeVar,
19 cast,
20)
21
22from hypothesis.errors import ChoiceTooLarge
23from hypothesis.internal.conjecture.floats import float_to_lex, lex_to_float
24from hypothesis.internal.conjecture.utils import identity
25from hypothesis.internal.floats import float_to_int, make_float_clamper, sign_aware_lte
26from hypothesis.internal.intervalsets import IntervalSet
27
28T = TypeVar("T")
29
30
31class IntegerConstraints(TypedDict):
32 min_value: int | None
33 max_value: int | None
34 weights: dict[int, float] | None
35 shrink_towards: int
36
37
38class FloatConstraints(TypedDict):
39 min_value: float
40 max_value: float
41 allow_nan: bool
42 smallest_nonzero_magnitude: float
43
44
45class StringConstraints(TypedDict):
46 intervals: IntervalSet
47 min_size: int
48 max_size: int
49
50
51class BytesConstraints(TypedDict):
52 min_size: int
53 max_size: int
54
55
56class BooleanConstraints(TypedDict):
57 p: float
58
59
60ChoiceT: TypeAlias = int | str | bool | float | bytes
61ChoiceConstraintsT: TypeAlias = (
62 IntegerConstraints
63 | FloatConstraints
64 | StringConstraints
65 | BytesConstraints
66 | BooleanConstraints
67)
68ChoiceTypeT: TypeAlias = Literal["integer", "string", "boolean", "float", "bytes"]
69ChoiceKeyT: TypeAlias = (
70 int | str | bytes | tuple[Literal["bool"], bool] | tuple[Literal["float"], int]
71)
72
73
74@dataclass(slots=True, frozen=False)
75class ChoiceTemplate:
76 type: Literal["simplest"]
77 count: int | None
78
79 def __post_init__(self) -> None:
80 if self.count is not None:
81 assert self.count > 0
82
83
84@dataclass(slots=True, frozen=True)
85class ValueHole:
86 """A hole in a choice sequence, carrying a value instead of choices.
87
88 When ``ConjectureData.draw`` finds a ValueHole at the current prefix
89 position, it asks the strategy being drawn to ``_invert`` the value, and
90 on success splices the resulting choices into the prefix in place of the
91 hole - so the value is re-encoded by whichever strategy is drawn at that
92 position, with that strategy's own constraints. If no strategy claims the
93 hole (inversion failed, or the hole fell out of alignment with strategy
94 draw boundaries), it is treated as a misalignment when a draw_* call
95 reaches it.
96 """
97
98 # any value a strategy might re-encode, not just the choice types
99 value: object
100
101
102@dataclass(slots=True, frozen=False)
103class ChoiceNode:
104 type: ChoiceTypeT
105 value: ChoiceT
106 constraints: ChoiceConstraintsT
107 was_forced: bool
108 index: int | None = None
109
110 def copy(
111 self,
112 *,
113 with_value: ChoiceT | None = None,
114 with_constraints: ChoiceConstraintsT | None = None,
115 ) -> "ChoiceNode":
116 # we may want to allow this combination in the future, but for now it's
117 # a footgun.
118 if self.was_forced:
119 assert with_value is None, "modifying a forced node doesn't make sense"
120 # explicitly not copying index. node indices are only assigned via
121 # SpanRecord. This prevents footguns with relying on stale indices
122 # after copying.
123 return ChoiceNode(
124 type=self.type,
125 value=self.value if with_value is None else with_value,
126 constraints=(
127 self.constraints if with_constraints is None else with_constraints
128 ),
129 was_forced=self.was_forced,
130 )
131
132 @property
133 def trivial(self) -> bool:
134 """
135 A node is trivial if it cannot be simplified any further. This does not
136 mean that modifying a trivial node can't produce simpler test cases when
137 viewing the tree as a whole. Just that when viewing this node in
138 isolation, this is the simplest the node can get.
139 """
140 if self.was_forced:
141 return True
142
143 if self.type != "float":
144 zero_value = choice_from_index(0, self.type, self.constraints)
145 return choice_equal(self.value, zero_value)
146 else:
147 constraints = cast(FloatConstraints, self.constraints)
148 min_value = constraints["min_value"]
149 max_value = constraints["max_value"]
150 shrink_towards = 0.0
151
152 if min_value == -math.inf and max_value == math.inf:
153 return choice_equal(self.value, shrink_towards)
154
155 if (
156 not math.isinf(min_value)
157 and not math.isinf(max_value)
158 and math.ceil(min_value) <= math.floor(max_value)
159 ):
160 # the interval contains an integer. the simplest integer is the
161 # one closest to shrink_towards
162 shrink_towards = max(math.ceil(min_value), shrink_towards)
163 shrink_towards = min(math.floor(max_value), shrink_towards)
164 return choice_equal(self.value, float(shrink_towards))
165
166 # the real answer here is "the value in [min_value, max_value] with
167 # the lowest denominator when represented as a fraction".
168 # It would be good to compute this correctly in the future, but it's
169 # also not incorrect to be conservative here.
170 return False
171
172 def __eq__(self, other: object) -> bool:
173 if not isinstance(other, ChoiceNode):
174 return NotImplemented
175
176 return (
177 self.type == other.type
178 and choice_equal(self.value, other.value)
179 and choice_constraints_equal(self.type, self.constraints, other.constraints)
180 and self.was_forced == other.was_forced
181 )
182
183 def __hash__(self) -> int:
184 return hash(
185 (
186 self.type,
187 choice_key(self.value),
188 choice_constraints_key(self.type, self.constraints),
189 self.was_forced,
190 )
191 )
192
193 def __repr__(self) -> str:
194 forced_marker = " [forced]" if self.was_forced else ""
195 return f"{self.type} {self.value!r}{forced_marker} {self.constraints!r}"
196
197
198def _size_to_index(size: int, *, alphabet_size: int) -> int:
199 # this is the closed form of this geometric series:
200 # for i in range(size):
201 # index += alphabet_size**i
202 if alphabet_size <= 0:
203 assert size == 0
204 return 0
205 if alphabet_size == 1:
206 return size
207 v = (alphabet_size**size - 1) // (alphabet_size - 1)
208 # mypy thinks (m: int) // (n: int) -> Any. assert it back to int.
209 return cast(int, v)
210
211
212def _index_to_size(index: int, alphabet_size: int) -> int:
213 if alphabet_size == 0:
214 return 0
215 elif alphabet_size == 1:
216 # there is only one string of each size, so the size is equal to its
217 # ordering.
218 return index
219
220 # the closed-form inverse of _size_to_index is
221 # size = math.floor(math.log(index * (alphabet_size - 1) + 1, alphabet_size))
222 # which is fast, but suffers from float precision errors. As performance is
223 # relatively critical here, we'll use this formula by default, but fall back to
224 # a much slower integer-only logarithm when the calculation is too close for
225 # comfort.
226 total = index * (alphabet_size - 1) + 1
227 size = math.log(total, alphabet_size)
228
229 # if this computation is close enough that it could have been affected by
230 # floating point errors, use a much slower integer-only logarithm instead,
231 # which is guaranteed to be precise.
232 if 0 < math.ceil(size) - size < 1e-7:
233 s = 0
234 while total >= alphabet_size:
235 total //= alphabet_size
236 s += 1
237 return s
238 return math.floor(size)
239
240
241def collection_index(
242 choice: Sequence[T],
243 *,
244 min_size: int,
245 alphabet_size: int,
246 to_order: Callable[[T], int],
247) -> int:
248 # Collections are ordered by counting the number of values of each size,
249 # starting with min_size. alphabet_size indicates how many options there
250 # are for a single element. to_order orders an element by returning an n ≥ 0.
251
252 # we start by adding the size to the index, relative to min_size.
253 index = _size_to_index(len(choice), alphabet_size=alphabet_size) - _size_to_index(
254 min_size, alphabet_size=alphabet_size
255 )
256 # We then add each element c to the index, starting from the end (so "ab" is
257 # simpler than "ba"). Each loop takes c at position i in the sequence and
258 # computes the number of sequences of size i which come before it in the ordering.
259
260 # this running_exp computation is equivalent to doing
261 # index += (alphabet_size**i) * n
262 # but reuses intermediate exponentiation steps for efficiency.
263 running_exp = 1
264 for c in reversed(choice):
265 index += running_exp * to_order(c)
266 running_exp *= alphabet_size
267 return index
268
269
270def collection_value(
271 index: int,
272 *,
273 min_size: int,
274 alphabet_size: int,
275 from_order: Callable[[int], T],
276) -> list[T]:
277 from hypothesis.internal.conjecture.engine import BUFFER_SIZE
278
279 # this function is probably easiest to make sense of as an inverse of
280 # collection_index, tracking ~corresponding lines of code between the two.
281
282 index += _size_to_index(min_size, alphabet_size=alphabet_size)
283 size = _index_to_size(index, alphabet_size=alphabet_size)
284 # index -> value computation can be arbitrarily expensive for arbitrarily
285 # large min_size collections. short-circuit if the resulting size would be
286 # obviously-too-large. callers will generally turn this into a .mark_overrun().
287 if size >= BUFFER_SIZE:
288 raise ChoiceTooLarge
289
290 # subtract out the amount responsible for the size
291 index -= _size_to_index(size, alphabet_size=alphabet_size)
292 vals: list[T] = []
293 for i in reversed(range(size)):
294 # optimization for common case when we hit index 0. Exponentiation
295 # on large integers is expensive!
296 if index == 0:
297 n = 0
298 else:
299 n = index // (alphabet_size**i)
300 # subtract out the nearest multiple of alphabet_size**i
301 index -= n * (alphabet_size**i)
302 vals.append(from_order(n))
303 return vals
304
305
306def zigzag_index(value: int, *, shrink_towards: int) -> int:
307 # value | 0 1 -1 2 -2 3 -3 4
308 # index | 0 1 2 3 4 5 6 7
309 index = 2 * abs(shrink_towards - value)
310 if value > shrink_towards:
311 index -= 1
312 return index
313
314
315def zigzag_value(index: int, *, shrink_towards: int) -> int:
316 assert index >= 0
317 # count how many "steps" away from shrink_towards we are.
318 n = (index + 1) // 2
319 # now check if we're stepping up or down from shrink_towards.
320 if (index % 2) == 0:
321 n *= -1
322 return shrink_towards + n
323
324
325def choice_to_index(choice: ChoiceT, constraints: ChoiceConstraintsT) -> int:
326 # This function takes a choice in the choice sequence and returns the
327 # complexity index of that choice from among its possible values, where 0
328 # is the simplest.
329 #
330 # Note that the index of a choice depends on its constraints. The simplest value
331 # (at index 0) for {"min_value": None, "max_value": None} is 0, while for
332 # {"min_value": 1, "max_value": None} the simplest value is 1.
333 #
334 # choice_from_index inverts this function. An invariant on both functions is
335 # that they must be injective. Unfortunately, floats do not currently respect
336 # this. That's not *good*, but nothing has blown up - yet. And ordering
337 # floats in a sane manner is quite hard, so I've left it for another day.
338
339 if isinstance(choice, int) and not isinstance(choice, bool):
340 # Let a = shrink_towards.
341 # * Unbounded: Ordered by (|a - x|, sgn(a - x)). Think of a zigzag.
342 # [a, a + 1, a - 1, a + 2, a - 2, ...]
343 # * Semi-bounded: Same as unbounded, except stop on one side when you hit
344 # {min, max}_value. so min_value=-1 a=0 has order
345 # [0, 1, -1, 2, 3, 4, ...]
346 # * Bounded: Same as unbounded and semibounded, except stop on each side
347 # when you hit {min, max}_value.
348 #
349 # To simplify and gain intuition about this ordering, you can think about
350 # the most common case where 0 is first (a = 0). We deviate from this only
351 # rarely, e.g. for datetimes, where we generally want year 2000 to be
352 # simpler than year 0.
353 constraints = cast(IntegerConstraints, constraints)
354 shrink_towards = constraints["shrink_towards"]
355 min_value = constraints["min_value"]
356 max_value = constraints["max_value"]
357
358 if min_value is not None:
359 shrink_towards = max(min_value, shrink_towards)
360 if max_value is not None:
361 shrink_towards = min(max_value, shrink_towards)
362
363 if min_value is None and max_value is None:
364 # case: unbounded
365 return zigzag_index(choice, shrink_towards=shrink_towards)
366 elif min_value is not None and max_value is None:
367 # case: semibounded below
368
369 # min_value = -2
370 # index | 0 1 2 3 4 5 6 7
371 # v | 0 1 -1 2 -2 3 4 5
372 if abs(choice - shrink_towards) <= (shrink_towards - min_value):
373 return zigzag_index(choice, shrink_towards=shrink_towards)
374 return choice - min_value
375 elif max_value is not None and min_value is None:
376 # case: semibounded above
377 if abs(choice - shrink_towards) <= (max_value - shrink_towards):
378 return zigzag_index(choice, shrink_towards=shrink_towards)
379 return max_value - choice
380 else:
381 # case: bounded
382
383 # range = [-2, 5]
384 # shrink_towards = 2
385 # index | 0 1 2 3 4 5 6 7
386 # v | 2 3 1 4 0 5 -1 -2
387 #
388 # ^ with zero weights at index = [0, 2, 6]
389 # index | 0 1 2 3 4
390 # v | 3 4 0 5 -2
391
392 assert min_value is not None
393 assert max_value is not None
394 assert constraints["weights"] is None or all(
395 w > 0 for w in constraints["weights"].values()
396 ), "technically possible but really annoying to support zero weights"
397
398 # check which side gets exhausted first
399 if (shrink_towards - min_value) < (max_value - shrink_towards):
400 # Below shrink_towards gets exhausted first. Equivalent to
401 # semibounded below
402 if abs(choice - shrink_towards) <= (shrink_towards - min_value):
403 return zigzag_index(choice, shrink_towards=shrink_towards)
404 return choice - min_value
405 else:
406 # Above shrink_towards gets exhausted first. Equivalent to semibounded
407 # above
408 if abs(choice - shrink_towards) <= (max_value - shrink_towards):
409 return zigzag_index(choice, shrink_towards=shrink_towards)
410 return max_value - choice
411 elif isinstance(choice, bool):
412 constraints = cast(BooleanConstraints, constraints)
413 # Ordered by [False, True].
414 p = constraints["p"]
415 if not (2 ** (-64) < p < (1 - 2 ** (-64))):
416 # only one option is possible, so whatever it is is first.
417 return 0
418 return int(choice)
419 elif isinstance(choice, bytes):
420 constraints = cast(BytesConstraints, constraints)
421 return collection_index(
422 list(choice),
423 min_size=constraints["min_size"],
424 alphabet_size=2**8,
425 to_order=identity,
426 )
427 elif isinstance(choice, str):
428 constraints = cast(StringConstraints, constraints)
429 intervals = constraints["intervals"]
430 return collection_index(
431 choice,
432 min_size=constraints["min_size"],
433 alphabet_size=len(intervals),
434 to_order=intervals.index_from_char_in_shrink_order,
435 )
436 elif isinstance(choice, float):
437 sign = int(math.copysign(1.0, choice) < 0)
438 return (sign << 64) | float_to_lex(abs(choice))
439 else:
440 raise NotImplementedError
441
442
443def choice_from_index(
444 index: int, choice_type: ChoiceTypeT, constraints: ChoiceConstraintsT
445) -> ChoiceT:
446 assert index >= 0
447 if choice_type == "integer":
448 constraints = cast(IntegerConstraints, constraints)
449 shrink_towards = constraints["shrink_towards"]
450 min_value = constraints["min_value"]
451 max_value = constraints["max_value"]
452
453 if min_value is not None:
454 shrink_towards = max(min_value, shrink_towards)
455 if max_value is not None:
456 shrink_towards = min(max_value, shrink_towards)
457
458 if min_value is None and max_value is None:
459 # case: unbounded
460 return zigzag_value(index, shrink_towards=shrink_towards)
461 elif min_value is not None and max_value is None:
462 # case: semibounded below
463 if index <= zigzag_index(min_value, shrink_towards=shrink_towards):
464 return zigzag_value(index, shrink_towards=shrink_towards)
465 return index + min_value
466 elif max_value is not None and min_value is None:
467 # case: semibounded above
468 if index <= zigzag_index(max_value, shrink_towards=shrink_towards):
469 return zigzag_value(index, shrink_towards=shrink_towards)
470 return max_value - index
471 else:
472 # case: bounded
473 assert min_value is not None
474 assert max_value is not None
475 assert constraints["weights"] is None or all(
476 w > 0 for w in constraints["weights"].values()
477 ), "possible but really annoying to support zero weights"
478
479 if (shrink_towards - min_value) < (max_value - shrink_towards):
480 # equivalent to semibounded below case
481 if index <= zigzag_index(min_value, shrink_towards=shrink_towards):
482 return zigzag_value(index, shrink_towards=shrink_towards)
483 return index + min_value
484 else:
485 # equivalent to semibounded above case
486 if index <= zigzag_index(max_value, shrink_towards=shrink_towards):
487 return zigzag_value(index, shrink_towards=shrink_towards)
488 return max_value - index
489 elif choice_type == "boolean":
490 constraints = cast(BooleanConstraints, constraints)
491 # Ordered by [False, True].
492 p = constraints["p"]
493 only = None
494 if p <= 2 ** (-64):
495 only = False
496 elif p >= (1 - 2 ** (-64)):
497 only = True
498
499 assert index in {0, 1}
500 if only is not None:
501 # only one choice
502 assert index == 0
503 return only
504 return bool(index)
505 elif choice_type == "bytes":
506 constraints = cast(BytesConstraints, constraints)
507 value_b = collection_value(
508 index,
509 min_size=constraints["min_size"],
510 alphabet_size=2**8,
511 from_order=identity,
512 )
513 return bytes(value_b)
514 elif choice_type == "string":
515 constraints = cast(StringConstraints, constraints)
516 intervals = constraints["intervals"]
517 # _s because mypy is unhappy with reusing different-typed names in branches,
518 # even if the branches are disjoint.
519 value_s = collection_value(
520 index,
521 min_size=constraints["min_size"],
522 alphabet_size=len(intervals),
523 from_order=intervals.char_in_shrink_order,
524 )
525 return "".join(value_s)
526 elif choice_type == "float":
527 constraints = cast(FloatConstraints, constraints)
528 sign = -1 if index >> 64 else 1
529 result = sign * lex_to_float(index & ((1 << 64) - 1))
530
531 clamper = make_float_clamper(
532 min_value=constraints["min_value"],
533 max_value=constraints["max_value"],
534 smallest_nonzero_magnitude=constraints["smallest_nonzero_magnitude"],
535 allow_nan=constraints["allow_nan"],
536 )
537 return clamper(result)
538 else:
539 raise NotImplementedError
540
541
542def choice_permitted(choice: ChoiceT, constraints: ChoiceConstraintsT) -> bool:
543 if isinstance(choice, int) and not isinstance(choice, bool):
544 constraints = cast(IntegerConstraints, constraints)
545 min_value = constraints["min_value"]
546 max_value = constraints["max_value"]
547 if min_value is not None and choice < min_value:
548 return False
549 return not (max_value is not None and choice > max_value)
550 elif isinstance(choice, float):
551 constraints = cast(FloatConstraints, constraints)
552 if math.isnan(choice):
553 return constraints["allow_nan"]
554 if 0 < abs(choice) < constraints["smallest_nonzero_magnitude"]:
555 return False
556 return sign_aware_lte(constraints["min_value"], choice) and sign_aware_lte(
557 choice, constraints["max_value"]
558 )
559 elif isinstance(choice, str):
560 constraints = cast(StringConstraints, constraints)
561 if len(choice) < constraints["min_size"]:
562 return False
563 if len(choice) > constraints["max_size"]:
564 return False
565 return all(ord(c) in constraints["intervals"] for c in choice)
566 elif isinstance(choice, bytes):
567 constraints = cast(BytesConstraints, constraints)
568 if len(choice) < constraints["min_size"]:
569 return False
570 return len(choice) <= constraints["max_size"]
571 elif isinstance(choice, bool):
572 constraints = cast(BooleanConstraints, constraints)
573 if constraints["p"] <= 0:
574 return choice is False
575 if constraints["p"] >= 1:
576 return choice is True
577 return True
578 else:
579 raise NotImplementedError(f"unhandled type {type(choice)} with value {choice}")
580
581
582def choices_key(choices: Sequence[ChoiceT]) -> tuple[ChoiceKeyT, ...]:
583 return tuple(choice_key(choice) for choice in choices)
584
585
586def choice_key(choice: ChoiceT) -> ChoiceKeyT:
587 if isinstance(choice, float):
588 # float_to_int to distinguish -0.0/0.0, signaling/nonsignaling nans, etc,
589 # and then add a "float" key to avoid colliding with actual integers.
590 return ("float", float_to_int(choice))
591 if isinstance(choice, bool):
592 # avoid choice_key(0) == choice_key(False)
593 return ("bool", choice)
594 return choice
595
596
597def choice_equal(choice1: ChoiceT, choice2: ChoiceT) -> bool:
598 assert type(choice1) is type(choice2), (choice1, choice2)
599 return choice_key(choice1) == choice_key(choice2)
600
601
602def choice_constraints_equal(
603 choice_type: ChoiceTypeT,
604 constraints1: ChoiceConstraintsT,
605 constraints2: ChoiceConstraintsT,
606) -> bool:
607 return choice_constraints_key(choice_type, constraints1) == choice_constraints_key(
608 choice_type, constraints2
609 )
610
611
612def choice_constraints_key(
613 choice_type: ChoiceTypeT, constraints: ChoiceConstraintsT
614) -> tuple[Hashable, ...]:
615 if choice_type == "float":
616 constraints = cast(FloatConstraints, constraints)
617 return (
618 float_to_int(constraints["min_value"]),
619 float_to_int(constraints["max_value"]),
620 constraints["allow_nan"],
621 constraints["smallest_nonzero_magnitude"],
622 )
623 if choice_type == "integer":
624 constraints = cast(IntegerConstraints, constraints)
625 return (
626 constraints["min_value"],
627 constraints["max_value"],
628 None if constraints["weights"] is None else tuple(constraints["weights"]),
629 constraints["shrink_towards"],
630 )
631 return tuple(constraints[key] for key in sorted(constraints)) # type: ignore
632
633
634def choices_size(choices: Iterable[ChoiceT]) -> int:
635 from hypothesis.database import choices_to_bytes
636
637 return len(choices_to_bytes(choices))