[pypy-commit] pypy stacklet: Implement switch(to=x).

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: stacklet
Changeset: r46606:55a35a42eeca
Date: 2011-08-18 16:39 +0200
http://bitbucket.org/pypy/pypy/changeset/55a35a42eeca/

Log:	Implement switch(to=x).

diff --git a/pypy/module/_continuation/interp_continuation.py b/pypy/module/_continuation/interp_continuation.py
--- a/pypy/module/_continuation/interp_continuation.py
+++ b/pypy/module/_continuation/interp_continuation.py
@@ -54,30 +54,27 @@
         if self.sthread.is_empty_handle(self.h):
             raise geterror(self.space, "continulet already finished")
         ec = self.check_sthread()
+        if self is to:    # double-switch to myself: no-op
+            return w_value
         saved_topframeref = ec.topframeref
         start_state.w_value = w_value
         #
+        start_state.origin = self
         if to is None:
-            start_state.origin = None
-            h = self.h      # simple switch: going to self.h
+            # simple switch: going to self.h
+            start_state.destination = self
         else:
-            start_state.origin = self
-            h = to.h        # double switch: the final destination is to.h
+            # double switch: the final destination is to.h
+            start_state.destination = to
         #
         try:
-            h = self.sthread.switch(h)
+            sthread = self.sthread
+            do_switch(sthread, start_state.destination.h)
         except MemoryError:
             start_state.clear()
             raise getmemoryerror(self.space)
         #
-        origin = start_state.origin
-        if origin is not None:
-            # double switch ('self' is now the final destination)
-            start_state.origin = None
-            h, origin.h = origin.h, h
-        self.h = h
-        #
-        ec = self.sthread.ec
+        ec = sthread.ec
         ec.topframeref = saved_topframeref
         if start_state.propagate_exception:
             e = start_state.propagate_exception
@@ -145,6 +142,7 @@
 class StartState:   # xxx a single global to pass around the function to start
     def clear(self):
         self.origin = None
+        self.destination = None
         self.w_callable = None
         self.args = None
         self.w_value = None
@@ -159,7 +157,7 @@
     args       = start_state.args
     start_state.clear()
     try:
-        self.h = self.sthread.switch(h)
+        do_switch(self.sthread, h)
     except MemoryError:
         return h       # oups!  do an early return in this case
     #
@@ -176,11 +174,21 @@
             raise OperationError(space.w_TypeError, space.wrap(
                 "can't send non-None value to a just-started continulet"))
 
-        args = args.prepend(space.wrap(self))
+        args = args.prepend(self.space.wrap(self))
         w_result = space.call_args(w_callable, args)
     except Exception, e:
         start_state.propagate_exception = e
-        return self.h
     else:
         start_state.w_value = w_result
-        return self.h
+    start_state.origin = self
+    start_state.destination = self
+    return self.h
+
+
+def do_switch(sthread, h):
+    h = sthread.switch(h)
+    origin = start_state.origin
+    self = start_state.destination
+    start_state.origin = None
+    start_state.destination = None
+    self.h, origin.h = origin.h, h
diff --git a/pypy/module/_continuation/test/test_stacklet.py b/pypy/module/_continuation/test/test_stacklet.py
--- a/pypy/module/_continuation/test/test_stacklet.py
+++ b/pypy/module/_continuation/test/test_stacklet.py
@@ -429,7 +429,7 @@
         c2 = continulet(f2)
         res = c1.switch()
         assert res == 42
-        assert not c2.is_pending()
+        assert not c2.is_pending()    # finished by returning 42
         res = c1.switch('c')
         assert res == 41
 
@@ -449,10 +449,12 @@
         from _continuation import continulet
         #
         def f1(c1):
+            print 'in f1'
             return 'm'
         #
         def f2(c2):
             res = c2.switch('z')
+            print 'got there!'
             assert res == 'a'
             return None
         #
@@ -460,7 +462,11 @@
         c2 = continulet(f2)
         res = c2.switch()
         assert res == 'z'
+        assert c1.is_pending()
+        assert c2.is_pending()
+        print 'calling!'
         res = c1.switch('a', to=c2)
+        print 'back'
         assert res == 'm'
 
     def test_switch2_immediately_away_corner_case(self):