Make it possible for rust functions to increase recursion depth (RustPython#3252)

Make it possible for rust functions to increase recursion depth

Author: pheki

Lib/test/list_tests.py

@@ -60,8 +60,6 @@ def test_repr(self):
         self.assertEqual(str(a2), "[0, 1, 2, [...], 3]")
         self.assertEqual(repr(a2), "[0, 1, 2, [...], 3]")
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_repr_deep(self):
         a = self.type2test([])
         for i in range(sys.getrecursionlimit() + 100):

Lib/test/test_dict.py

@@ -545,8 +545,6 @@ def __repr__(self):
         d = {1: BadRepr()}
         self.assertRaises(Exc, repr, d)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @unittest.skipIf(sys.platform == 'win32', 'TODO: RUSTPYTHON Windows')
     def test_repr_deep(self):
         d = {}

Lib/test/test_dictviews.py

@@ -224,8 +224,6 @@ def test_recursive_repr(self):
         # Again.
         self.assertIsInstance(r, str)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @unittest.skipIf(sys.platform == "win32", "thread 'main' has overflowed its stack")
     def test_deeply_nested_repr(self):
         d = {}

extra_tests/snippets/recursion.py

@@ -0,0 +1,11 @@
+from testutils import assert_raises
+
+class Foo(object):
+    pass
+
+Foo.__repr__ = Foo.__str__
+
+foo = Foo()
+# Since the default __str__ implementation calls __repr__ and __repr__ is
+# actually __str__, str(foo) should raise a RecursionError.
+assert_raises(RecursionError, str, foo)

vm/src/stdlib/sys.rs

@@ -478,16 +478,16 @@ mod sys {
                     "recursion limit must be greater than or equal to one".to_owned(),
                 )
             })?;
-        let recursion_depth = vm.frames.borrow().len();
+        let recursion_depth = vm.current_recursion_depth();
 
         if recursion_limit > recursion_depth + 1 {
             vm.recursion_limit.set(recursion_limit);
             Ok(())
         } else {
             Err(vm.new_recursion_error(format!(
-            "cannot set the recursion limit to {} at the recursion depth {}: the limit is too low",
-            recursion_limit, recursion_depth
-        )))
+                "cannot set the recursion limit to {} at the recursion depth {}: the limit is too low",
+                recursion_limit, recursion_depth
+            )))
         }
     }
 

vm/src/vm.rs

@@ -62,6 +62,7 @@ pub struct VirtualMachine {
     pub repr_guards: RefCell<HashSet<usize>>,
     pub state: PyRc<PyGlobalState>,
     pub initialized: bool,
+    recursion_depth: Cell<usize>,
 }
 
 #[derive(Debug, Default)]
@@ -297,6 +298,7 @@ impl VirtualMachine {
                 codec_registry,
             }),
             initialized: false,
+            recursion_depth: Cell::new(0),
         };
 
         let frozen = frozen::map_frozen(&vm, frozen::get_module_inits()).collect();
@@ -464,6 +466,7 @@ impl VirtualMachine {
             repr_guards: RefCell::default(),
             state: self.state.clone(),
             initialized: self.initialized,
+            recursion_depth: Cell::new(0),
         };
         PyThread { thread_vm }
     }
@@ -502,25 +505,41 @@ impl VirtualMachine {
         }
     }
 
+    pub fn current_recursion_depth(&self) -> usize {
+        self.recursion_depth.get()
+    }
+
+    /// Used to run the body of a (possibly) recursive function. It will raise a
+    /// RecursionError if recursive functions are nested far too many times,
+    /// preventing a stack overflow.
+    pub fn with_recursion<R, F: FnOnce() -> PyResult<R>>(&self, _where: &str, f: F) -> PyResult<R> {
+        self.check_recursive_call(_where)?;
+        self.recursion_depth.set(self.recursion_depth.get() + 1);
+        let result = f();
+        self.recursion_depth.set(self.recursion_depth.get() - 1);
+        result
+    }
+
     pub fn with_frame<R, F: FnOnce(FrameRef) -> PyResult<R>>(
         &self,
         frame: FrameRef,
         f: F,
     ) -> PyResult<R> {
-        self.check_recursive_call("")?;
-        self.frames.borrow_mut().push(frame.clone());
-        let result = f(frame);
-        // defer dec frame
-        let _popped = self.frames.borrow_mut().pop();
-        result
+        self.with_recursion("", || {
+            self.frames.borrow_mut().push(frame.clone());
+            let result = f(frame);
+            // defer dec frame
+            let _popped = self.frames.borrow_mut().pop();
+            result
+        })
     }
 
     pub fn run_frame(&self, frame: FrameRef) -> PyResult<ExecutionResult> {
         self.with_frame(frame, |f| f.run(self))
     }
 
     fn check_recursive_call(&self, _where: &str) -> PyResult<()> {
-        if self.frames.borrow().len() > self.recursion_limit.get() {
+        if self.recursion_depth.get() > self.recursion_limit.get() {
             Err(self.new_recursion_error(format!("maximum recursion depth exceeded {}", _where)))
         } else {
             Ok(())
@@ -887,8 +906,10 @@ impl VirtualMachine {
     }
 
     pub fn to_repr(&self, obj: &PyObjectRef) -> PyResult<PyStrRef> {
-        let repr = self.call_special_method(obj.clone(), "__repr__", ())?;
-        repr.try_into_value(self)
+        self.with_recursion(" while getting the repr of an object", || {
+            let repr = self.call_special_method(obj.clone(), "__repr__", ())?;
+            repr.try_into_value(self)
+        })
     }
 
     pub fn to_index_opt(&self, obj: PyObjectRef) -> Option<PyResult<PyIntRef>> {