Skip to content

Commit 66e9d81

Browse files
committed
handle case where a string must be freed and copied from and to the same arena
1 parent 0c9e400 commit 66e9d81

File tree

1 file changed

+33
-5
lines changed

1 file changed

+33
-5
lines changed

stringdtype/stringdtype/src/static_string.c

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,12 @@ is_not_a_vstring(const npy_packed_static_string *s)
215215
return is_short_string(s) || npy_string_isnull(s);
216216
}
217217

218+
int
219+
is_a_vstring(const npy_packed_static_string *s)
220+
{
221+
return !is_not_a_vstring(s);
222+
}
223+
218224
int
219225
npy_string_load(npy_string_allocator *allocator,
220226
const npy_packed_static_string *packed_string,
@@ -281,7 +287,7 @@ heap_or_arena_allocate(npy_string_allocator *allocator,
281287
return buf;
282288
}
283289
else {
284-
// no room, resort to a heap allocation this leaves the
290+
// No room, resort to a heap allocation. This leaves the
285291
// NPY_STRING_ARENA_FREED flag set to possibly re-use the arena
286292
// allocation in the future if there is room for it
287293
*flags |= NPY_STRING_ON_HEAP;
@@ -454,14 +460,36 @@ npy_string_dup(const npy_packed_static_string *in,
454460
memcpy(out, in, sizeof(npy_packed_static_string));
455461
return 0;
456462
}
457-
458463
_npy_static_string_u *in_u = (_npy_static_string_u *)in;
464+
size_t size = VSTRING_SIZE(in_u);
465+
if (size == 0) {
466+
_npy_static_string_u *out_u = (_npy_static_string_u *)out;
467+
unsigned char flags = out_u->direct_buffer.flags_and_size &
468+
~NPY_SHORT_STRING_SIZE_MASK;
469+
*out = *NPY_EMPTY_STRING;
470+
out_u->direct_buffer.flags_and_size |= flags;
471+
return 0;
472+
}
473+
char *in_buf = NULL;
459474
npy_string_arena *arena = &in_allocator->arena;
460-
if (arena == NULL) {
475+
if (arena->buffer == NULL) {
461476
return -1;
462477
}
463-
return npy_string_newsize(vstring_buffer(arena, in_u), VSTRING_SIZE(in_u),
464-
out, out_allocator);
478+
int used_malloc = 0;
479+
if (in_allocator == out_allocator && is_a_vstring(in)) {
480+
in_buf = in_allocator->malloc(size);
481+
memcpy(in_buf, vstring_buffer(arena, in_u), size);
482+
used_malloc = 1;
483+
}
484+
else {
485+
in_buf = vstring_buffer(arena, in_u);
486+
}
487+
int ret =
488+
npy_string_newsize(in_buf, VSTRING_SIZE(in_u), out, out_allocator);
489+
if (used_malloc) {
490+
in_allocator->free(in_buf);
491+
}
492+
return ret;
465493
}
466494

467495
int

0 commit comments

Comments
 (0)