Finish() shrinks the frozenset if needed

Author: vstinner

Lib/test/test_capi/test_set.py

@@ -296,14 +296,19 @@ def assert_unhashable(self):
         return self.assertRaisesRegex(TypeError, err_msg)
 
     def test_empty(self):
-        for size in (0, 128):
+        _set_get_table_size = _testcapi._set_get_table_size
+
+        for size in (0, 1024):
             with self.subTest(size=size):
                 writer = self.create_writer(size)
                 empty_tuple = ()
                 writer.update(empty_tuple)
                 writer.update_array(empty_tuple)
                 writer.update_array_steal(empty_tuple)
-                self.assertEqual(writer.finish(), frozenset())
+                result = writer.finish()
+                self.assertEqual(result, frozenset())
+                self.assertEqual(_set_get_table_size(result),
+                                 self.PySet_MINSIZE)
 
     def test_get_table_set(self):
         # Test size parameter of PyFrozenSetWriter_Create()
@@ -420,6 +425,24 @@ class MyObject:
         writer.update([obj])
         self.assertTrue(gc.is_tracked(writer.finish()))
 
+    def test_finish_shrink(self):
+        # Check that PyFrozenSetWriter_Finish() shrinks the table
+        # if the frozenset table is way too big
+        _set_get_table_size = _testcapi._set_get_table_size
+
+        writer = self.create_writer(1024)
+        writer.update([1])
+        result = writer.finish()
+        self.assertEqual(result, frozenset([1]))
+        self.assertEqual(_set_get_table_size(result), self.PySet_MINSIZE)
+
+        items = self.create_items(127)
+        writer = self.create_writer(len(items) * 100)
+        writer.update(items)
+        result = writer.finish()
+        self.assertEqual(result, frozenset(items))
+        self.assertEqual(_set_get_table_size(result), 256)
+
     def test_discard(self):
         # Test PyFrozenSetWriter_Discard(): must not leak memory
 

Modules/_testcapi/set.c

@@ -31,6 +31,15 @@ test_set_type_size(PyObject *self, PyObject *Py_UNUSED(ignored))
 }
 
 
+static PyObject*
+_set_get_table_size(PyObject *self_raw, PyObject *arg)
+{
+    PySetObject *frozenset = _PySet_CAST(arg);
+    Py_ssize_t table_size = frozenset->mask + 1;
+    return PyLong_FromSsize_t(table_size);
+}
+
+
 typedef struct {
     PyObject_HEAD
     PyFrozenSetWriter *writer;
@@ -210,9 +219,8 @@ writer_get_table_size(PyObject *self_raw, PyObject *Py_UNUSED(args))
         return NULL;
     }
 
-    PySetObject *frozenset = _PySet_CAST(self->writer);
-    Py_ssize_t table_size = frozenset->mask + 1;
-    return PyLong_FromSsize_t(table_size);
+    PyObject *frozenset = _PyObject_CAST(self->writer);
+    return _set_get_table_size(self_raw, frozenset);
 }
 
 
@@ -277,6 +285,7 @@ static PyType_Spec Writer_spec = {
 static PyMethodDef test_methods[] = {
     {"set_get_size", set_get_size, METH_O},
     {"test_set_type_size", test_set_type_size, METH_NOARGS},
+    {"_set_get_table_size", _set_get_table_size, METH_O},
     {NULL},
 };
 

Objects/setobject.c

@@ -3289,6 +3289,22 @@ PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer)
     assert(frozensetwriter_check_consistency(writer));
     PySetObject *frozenset = _PySet_CAST(writer);
 
+    // If the frozenset table is too big (waste memory), make the table
+    // smaller. For example, PyFrozenSetWriter_Create() argument was way too
+    // big.
+    Py_ssize_t table_size = frozenset->mask + 1;
+    // set_add_entry_takeref() allocates a table of used*4 for table smaller
+    // than 50,000 items. So check for table_size / 8.
+    if (frozenset->used < table_size / 8
+        && frozenset->table != frozenset->smalltable)
+    {
+        if (set_table_resize(frozenset, frozenset->used * 2) != 0) {
+            PyFrozenSetWriter_Discard(writer);
+            return NULL;
+        }
+        assert((frozenset->mask + 1) != table_size);
+    }
+
     // Now that the frozenset is fully initialized, it's safe to track it in
     // the GC. Track the frozenset if at least one item is tracked by the GC.
     Py_ssize_t pos = 0;