1 #include "Python.h"
2 #include "pycore_code.h" // _PyCode_GetVarnames()
3 #include "pycore_frame.h"
4 #include "pycore_pyerrors.h" // export _Py_UTF8_Edit_Cost()
5 #include "pycore_runtime.h" // _Py_ID()
6
7 #define MAX_CANDIDATE_ITEMS 750
8 #define MAX_STRING_SIZE 40
9
10 #define MOVE_COST 2
11 #define CASE_COST 1
12
13 #define LEAST_FIVE_BITS(n) ((n) & 31)
14
15 static inline int
substitution_cost(char a,char b)16 substitution_cost(char a, char b)
17 {
18 if (LEAST_FIVE_BITS(a) != LEAST_FIVE_BITS(b)) {
19 // Not the same, not a case flip.
20 return MOVE_COST;
21 }
22 if (a == b) {
23 return 0;
24 }
25 if ('A' <= a && a <= 'Z') {
26 a += ('a' - 'A');
27 }
28 if ('A' <= b && b <= 'Z') {
29 b += ('a' - 'A');
30 }
31 if (a == b) {
32 return CASE_COST;
33 }
34 return MOVE_COST;
35 }
36
37 /* Calculate the Levenshtein distance between string1 and string2 */
38 static Py_ssize_t
levenshtein_distance(const char * a,size_t a_size,const char * b,size_t b_size,size_t max_cost,size_t * buffer)39 levenshtein_distance(const char *a, size_t a_size,
40 const char *b, size_t b_size,
41 size_t max_cost, size_t *buffer)
42 {
43 // Both strings are the same (by identity)
44 if (a == b) {
45 return 0;
46 }
47
48 // Trim away common affixes.
49 while (a_size && b_size && a[0] == b[0]) {
50 a++; a_size--;
51 b++; b_size--;
52 }
53 while (a_size && b_size && a[a_size-1] == b[b_size-1]) {
54 a_size--;
55 b_size--;
56 }
57 if (a_size == 0 || b_size == 0) {
58 return (a_size + b_size) * MOVE_COST;
59 }
60 if (a_size > MAX_STRING_SIZE || b_size > MAX_STRING_SIZE) {
61 return max_cost + 1;
62 }
63
64 // Prefer shorter buffer
65 if (b_size < a_size) {
66 const char *t = a; a = b; b = t;
67 size_t t_size = a_size; a_size = b_size; b_size = t_size;
68 }
69
70 // quick fail when a match is impossible.
71 if ((b_size - a_size) * MOVE_COST > max_cost) {
72 return max_cost + 1;
73 }
74
75 // Instead of producing the whole traditional len(a)-by-len(b)
76 // matrix, we can update just one row in place.
77 // Initialize the buffer row
78 size_t tmp = MOVE_COST;
79 for (size_t i = 0; i < a_size; i++) {
80 // cost from b[:0] to a[:i+1]
81 buffer[i] = tmp;
82 tmp += MOVE_COST;
83 }
84
85 size_t result = 0;
86 for (size_t b_index = 0; b_index < b_size; b_index++) {
87 char code = b[b_index];
88 // cost(b[:b_index], a[:0]) == b_index * MOVE_COST
89 size_t distance = result = b_index * MOVE_COST;
90 size_t minimum = SIZE_MAX;
91 for (size_t index = 0; index < a_size; index++) {
92
93 // cost(b[:b_index+1], a[:index+1]) = min(
94 // // 1) substitute
95 // cost(b[:b_index], a[:index])
96 // + substitution_cost(b[b_index], a[index]),
97 // // 2) delete from b
98 // cost(b[:b_index], a[:index+1]) + MOVE_COST,
99 // // 3) delete from a
100 // cost(b[:b_index+1], a[index]) + MOVE_COST
101 // )
102
103 // 1) Previous distance in this row is cost(b[:b_index], a[:index])
104 size_t substitute = distance + substitution_cost(code, a[index]);
105 // 2) cost(b[:b_index], a[:index+1]) from previous row
106 distance = buffer[index];
107 // 3) existing result is cost(b[:b_index+1], a[index])
108
109 size_t insert_delete = Py_MIN(result, distance) + MOVE_COST;
110 result = Py_MIN(insert_delete, substitute);
111
112 // cost(b[:b_index+1], a[:index+1])
113 buffer[index] = result;
114 if (result < minimum) {
115 minimum = result;
116 }
117 }
118 if (minimum > max_cost) {
119 // Everything in this row is too big, so bail early.
120 return max_cost + 1;
121 }
122 }
123 return result;
124 }
125
126 PyObject *
_Py_CalculateSuggestions(PyObject * dir,PyObject * name)127 _Py_CalculateSuggestions(PyObject *dir,
128 PyObject *name)
129 {
130 assert(!PyErr_Occurred());
131 assert(PyList_CheckExact(dir));
132
133 Py_ssize_t dir_size = PyList_GET_SIZE(dir);
134 if (dir_size >= MAX_CANDIDATE_ITEMS) {
135 return NULL;
136 }
137
138 Py_ssize_t suggestion_distance = PY_SSIZE_T_MAX;
139 PyObject *suggestion = NULL;
140 Py_ssize_t name_size;
141 const char *name_str = PyUnicode_AsUTF8AndSize(name, &name_size);
142 if (name_str == NULL) {
143 return NULL;
144 }
145 size_t *buffer = PyMem_New(size_t, MAX_STRING_SIZE);
146 if (buffer == NULL) {
147 return PyErr_NoMemory();
148 }
149 for (int i = 0; i < dir_size; ++i) {
150 PyObject *item = PyList_GET_ITEM(dir, i);
151 if (_PyUnicode_Equal(name, item)) {
152 continue;
153 }
154 Py_ssize_t item_size;
155 const char *item_str = PyUnicode_AsUTF8AndSize(item, &item_size);
156 if (item_str == NULL) {
157 PyMem_Free(buffer);
158 return NULL;
159 }
160 // No more than 1/3 of the involved characters should need changed.
161 Py_ssize_t max_distance = (name_size + item_size + 3) * MOVE_COST / 6;
162 // Don't take matches we've already beaten.
163 max_distance = Py_MIN(max_distance, suggestion_distance - 1);
164 Py_ssize_t current_distance =
165 levenshtein_distance(name_str, name_size, item_str,
166 item_size, max_distance, buffer);
167 if (current_distance > max_distance) {
168 continue;
169 }
170 if (!suggestion || current_distance < suggestion_distance) {
171 suggestion = item;
172 suggestion_distance = current_distance;
173 }
174 }
175 PyMem_Free(buffer);
176 return Py_XNewRef(suggestion);
177 }
178
179 Py_ssize_t
_Py_UTF8_Edit_Cost(PyObject * a,PyObject * b,Py_ssize_t max_cost)180 _Py_UTF8_Edit_Cost(PyObject *a, PyObject *b, Py_ssize_t max_cost)
181 {
182 assert(PyUnicode_Check(a) && PyUnicode_Check(b));
183 Py_ssize_t size_a, size_b;
184 const char *utf8_a = PyUnicode_AsUTF8AndSize(a, &size_a);
185 if (utf8_a == NULL) {
186 return -1;
187 }
188 const char *utf8_b = PyUnicode_AsUTF8AndSize(b, &size_b);
189 if (utf8_b == NULL) {
190 return -1;
191 }
192 if (max_cost == -1) {
193 max_cost = MOVE_COST * Py_MAX(size_a, size_b);
194 }
195 size_t *buffer = PyMem_New(size_t, MAX_STRING_SIZE);
196 if (buffer == NULL) {
197 PyErr_NoMemory();
198 return -1;
199 }
200 Py_ssize_t res = levenshtein_distance(utf8_a, size_a,
201 utf8_b, size_b, max_cost, buffer);
202 PyMem_Free(buffer);
203 return res;
204 }
205
206