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
11from collections import defaultdict
12from random import Random
13from typing import Callable, Dict, Iterable, List, Optional, Sequence
14
15from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
16
17
18def prefix_selection_order(
19 prefix: Sequence[int],
20) -> Callable[[int, int], Iterable[int]]:
21 """Select choices starting from ``prefix```,
22 preferring to move left then wrapping around
23 to the right."""
24
25 def selection_order(depth: int, n: int) -> Iterable[int]:
26 if depth < len(prefix):
27 i = prefix[depth]
28 if i >= n:
29 i = n - 1
30 yield from range(i, -1, -1)
31 yield from range(n - 1, i, -1)
32 else:
33 yield from range(n - 1, -1, -1)
34
35 return selection_order
36
37
38def random_selection_order(random: Random) -> Callable[[int, int], Iterable[int]]:
39 """Select choices uniformly at random."""
40
41 def selection_order(depth: int, n: int) -> Iterable[int]:
42 pending = LazySequenceCopy(range(n))
43 while pending:
44 i = random.randrange(0, len(pending))
45 yield pending.pop(i)
46
47 return selection_order
48
49
50class Chooser:
51 """A source of nondeterminism for use in shrink passes."""
52
53 def __init__(
54 self,
55 tree: "ChoiceTree",
56 selection_order: Callable[[int, int], Iterable[int]],
57 ):
58 self.__selection_order = selection_order
59 self.__node_trail = [tree.root]
60 self.__choices: "List[int]" = []
61 self.__finished = False
62
63 def choose(
64 self,
65 values: Sequence[int],
66 condition: Callable[[int], bool] = lambda x: True,
67 ) -> int:
68 """Return some element of values satisfying the condition
69 that will not lead to an exhausted branch, or raise DeadBranch
70 if no such element exist".
71 """
72 assert not self.__finished
73 node = self.__node_trail[-1]
74 if node.live_child_count is None:
75 node.live_child_count = len(values)
76 node.n = len(values)
77
78 assert node.live_child_count > 0 or len(values) == 0
79
80 for i in self.__selection_order(len(self.__choices), len(values)):
81 if node.live_child_count == 0:
82 break
83 if not node.children[i].exhausted:
84 v = values[i]
85 if condition(v):
86 self.__choices.append(i)
87 self.__node_trail.append(node.children[i])
88 return v
89 else:
90 node.children[i] = DeadNode
91 node.live_child_count -= 1
92 assert node.live_child_count == 0
93 raise DeadBranch
94
95 def finish(self) -> Sequence[int]:
96 """Record the decisions made in the underlying tree and return
97 a prefix that can be used for the next Chooser to be used."""
98 self.__finished = True
99 assert len(self.__node_trail) == len(self.__choices) + 1
100
101 result = tuple(self.__choices)
102
103 self.__node_trail[-1].live_child_count = 0
104 while len(self.__node_trail) > 1 and self.__node_trail[-1].exhausted:
105 self.__node_trail.pop()
106 assert len(self.__node_trail) == len(self.__choices)
107 i = self.__choices.pop()
108 target = self.__node_trail[-1]
109 target.children[i] = DeadNode
110 assert target.live_child_count is not None
111 target.live_child_count -= 1
112
113 return result
114
115
116class ChoiceTree:
117 """Records sequences of choices made during shrinking so that we
118 can track what parts of a pass has run. Used to create Chooser
119 objects that are the main interface that a pass uses to make
120 decisions about what to do.
121 """
122
123 def __init__(self) -> None:
124 self.root = TreeNode()
125
126 @property
127 def exhausted(self) -> bool:
128 return self.root.exhausted
129
130 def step(
131 self,
132 selection_order: Callable[[int, int], Iterable[int]],
133 f: Callable[[Chooser], None],
134 ) -> Sequence[int]:
135 assert not self.exhausted
136
137 chooser = Chooser(self, selection_order)
138 try:
139 f(chooser)
140 except DeadBranch:
141 pass
142 return chooser.finish()
143
144
145class TreeNode:
146 def __init__(self) -> None:
147 self.children: Dict[int, TreeNode] = defaultdict(TreeNode)
148 self.live_child_count: "Optional[int]" = None
149 self.n: "Optional[int]" = None
150
151 @property
152 def exhausted(self) -> bool:
153 return self.live_child_count == 0
154
155
156DeadNode = TreeNode()
157DeadNode.live_child_count = 0
158
159
160class DeadBranch(Exception):
161 pass