diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index f877566c17e2..6743935b96da 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -265,8 +265,15 @@ VecNested VecNested_Append(VecNested vec, VecNestedBufItem x) { memcpy(new.buf->items, vec.buf->items, sizeof(VecNestedBufItem) * vec.len); // TODO: How to safely represent deleted items? memset(new.buf->items + vec.len, 0, sizeof(VecNestedBufItem) * (new_size - vec.len)); - // Clear the items in the old vec. We avoid reference count manipulation. - memset(vec.buf->items, 0, sizeof(VecNestedBufItem) * vec.len); + if (Py_REFCNT(vec.buf) > 1) { + // Other references to old buffer exist; INCREF items in new buffer + // so old buffer keeps valid references for aliases. + for (Py_ssize_t i = 0; i < vec.len; i++) + Py_XINCREF(new.buf->items[i].buf); + } else { + // No aliases; transfer ownership by clearing old buffer items. + memset(vec.buf->items, 0, sizeof(VecNestedBufItem) * vec.len); + } new.buf->items[vec.len] = x; new.len = vec.len + 1; VEC_DECREF(vec); diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index 56188a42410e..a68712d69b7b 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/mypyc/lib-rt/vecs/vec_t.c @@ -279,8 +279,15 @@ VecT VecT_Append(VecT vec, PyObject *x, size_t item_type) { // Copy items to new vec. memcpy(new.buf->items, vec.buf->items, sizeof(PyObject *) * vec.len); memset(new.buf->items + vec.len, 0, sizeof(PyObject *) * (new_size - vec.len)); - // Clear the items in the old vec. We avoid reference count manipulation. - memset(vec.buf->items, 0, sizeof(PyObject *) * vec.len); + if (Py_REFCNT(vec.buf) > 1) { + // Other references to old buffer exist; INCREF items in new buffer + // so old buffer keeps valid references for aliases. + for (Py_ssize_t i = 0; i < vec.len; i++) + Py_XINCREF(new.buf->items[i]); + } else { + // No aliases; transfer ownership by clearing old buffer items. + memset(vec.buf->items, 0, sizeof(PyObject *) * vec.len); + } new.buf->items[vec.len] = x; new.len = vec.len + 1; VEC_DECREF(vec); diff --git a/mypyc/test-data/run-vecs-nested-interp.test b/mypyc/test-data/run-vecs-nested-interp.test index 50f8bca600b9..d4b0af15aa18 100644 --- a/mypyc/test-data/run-vecs-nested-interp.test +++ b/mypyc/test-data/run-vecs-nested-interp.test @@ -455,3 +455,15 @@ def test_contains_deeply_nested() -> None: assert vec[vec[str]]([vec[str](['a'])]) in v assert vec[vec[str]]([vec[str](['b'])]) not in v assert vec[vec[str]]() not in v + +def test_append_with_alias() -> None: + v = vec[vec[i64]]() + v = append(v, vec[i64]([1, 2, 3])) + alias = v + v = append(v, vec[i64]([4, 5])) + v = append(v, vec[i64]([6])) + v = append(v, vec[i64]([7])) + assert alias[0] == vec[i64]([1, 2, 3]) + assert len(alias) == 1 + assert v[0] == vec[i64]([1, 2, 3]) + assert len(v) == 4 diff --git a/mypyc/test-data/run-vecs-nested.test b/mypyc/test-data/run-vecs-nested.test index 0cf430d9c06c..49a6b892f20b 100644 --- a/mypyc/test-data/run-vecs-nested.test +++ b/mypyc/test-data/run-vecs-nested.test @@ -506,3 +506,15 @@ def test_union_of_nested_vecs() -> None: with assertRaises(TypeError): a5: Any = vec[vec[vec[i64]]]() len_union(a5) + +def test_append_with_alias() -> None: + v = vec[vec[i64]]() + v = append(v, vec[i64]([1, 2, 3])) + alias = v + v = append(v, vec[i64]([4, 5])) + v = append(v, vec[i64]([6])) + v = append(v, vec[i64]([7])) + assert alias[0] == vec[i64]([1, 2, 3]) + assert len(alias) == 1 + assert v[0] == vec[i64]([1, 2, 3]) + assert len(v) == 4 diff --git a/mypyc/test-data/run-vecs-t-interp.test b/mypyc/test-data/run-vecs-t-interp.test index 76bb57bfdd57..aab522acc9ce 100644 --- a/mypyc/test-data/run-vecs-t-interp.test +++ b/mypyc/test-data/run-vecs-t-interp.test @@ -533,6 +533,18 @@ def test_pop_index() -> None: assert item == '15' assert v == vec[str]() +def test_append_with_alias() -> None: + v = vec[str]() + v = append(v, 'a') + alias = v + v = append(v, 'b') + v = append(v, 'c') + v = append(v, 'd') + assert alias[0] == 'a' + assert len(alias) == 1 + assert v[0] == 'a' + assert len(v) == 4 + [file helpers.py] # Test helper classes diff --git a/mypyc/test-data/run-vecs-t.test b/mypyc/test-data/run-vecs-t.test index 2b69e3140820..6855aba68d57 100644 --- a/mypyc/test-data/run-vecs-t.test +++ b/mypyc/test-data/run-vecs-t.test @@ -436,3 +436,15 @@ def test_union_of_vecs() -> None: with assertRaises(TypeError): a4: Any = vec[Foo]() len_union(a4) + +def test_append_with_alias() -> None: + v = vec[str]() + v = append(v, 'a') + alias = v + v = append(v, 'b') + v = append(v, 'c') + v = append(v, 'd') + assert alias[0] == 'a' + assert len(alias) == 1 + assert v[0] == 'a' + assert len(v) == 4