gh-128961: fix array iterator segfault on __setstate__

Author: tom-pytel

Lib/test/test_array.py

@@ -1665,5 +1665,11 @@ def test_tolist(self, size):
         self.assertEqual(ls[:8], list(example[:8]))
         self.assertEqual(ls[-8:], list(example[-8:]))
 
+    def test_gh_128961(self):
+        a = array.array('i')
+        it = iter(a)
+        list(it)
+        it.__setstate__(0)
+
 if __name__ == "__main__":
     unittest.main()

Modules/arraymodule.c

@@ -3090,11 +3090,14 @@ array_arrayiterator___setstate__(arrayiterobject *self, PyObject *state)
     Py_ssize_t index = PyLong_AsSsize_t(state);
     if (index == -1 && PyErr_Occurred())
         return NULL;
-    if (index < 0)
-        index = 0;
-    else if (index > Py_SIZE(self->ao))
-        index = Py_SIZE(self->ao); /* iterator exhausted */
-    self->index = index;
+    arrayobject *ao = self->ao;
+    // if (ao != NULL) {
+        if (index < 0)
+            index = 0;
+        else if (index > Py_SIZE(ao))
+            index = Py_SIZE(ao); /* iterator exhausted */
+        self->index = index;
+    // }
     Py_RETURN_NONE;
 }