[pypy-commit] pypy default: Reduce the number of cases of shadowstack root walking:

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r51927:256d3279bde5
Date: 2012-01-28 18:20 +0100
http://bitbucket.org/pypy/pypy/changeset/256d3279bde5/

Log:	Reduce the number of cases of shadowstack root walking: from 4
	(jit/nojit, thread/nothread) to 2 (thread/nothread).

diff --git a/pypy/jit/backend/llsupport/gc.py b/pypy/jit/backend/llsupport/gc.py
--- a/pypy/jit/backend/llsupport/gc.py
+++ b/pypy/jit/backend/llsupport/gc.py
@@ -433,6 +433,9 @@
         class RootIterator:
             _alloc_flavor_ = "raw"
 
+            def setcontext(iself, context):
+                iself.context = context
+
             def next(iself, gc, next, range_highest):
                 # Return the "next" valid GC object' address.  This usually
                 # means just returning "next", until we reach "range_highest",
diff --git a/pypy/rpython/memory/gctransform/shadowstack.py b/pypy/rpython/memory/gctransform/shadowstack.py
--- a/pypy/rpython/memory/gctransform/shadowstack.py
+++ b/pypy/rpython/memory/gctransform/shadowstack.py
@@ -27,28 +27,17 @@
             return top
         self.decr_stack = decr_stack
 
-        translator = gctransformer.translator
-        if (hasattr(translator, '_jit2gc') and
-                'root_iterator' in translator._jit2gc):
-            root_iterator = translator._jit2gc['root_iterator']
-            def jit_walk_stack_root(callback, addr, end):
-                root_iterator.context = NonConstant(llmemory.NULL)
-                gc = self.gc
-                while True:
-                    addr = root_iterator.next(gc, addr, end)
-                    if addr == llmemory.NULL:
-                        return
-                    callback(gc, addr)
-                    addr += sizeofaddr
-            self.rootstackhook = jit_walk_stack_root
-        else:
-            def default_walk_stack_root(callback, addr, end):
-                gc = self.gc
-                while addr != end:
-                    if gc.points_to_valid_gc_object(addr):
-                        callback(gc, addr)
-                    addr += sizeofaddr
-            self.rootstackhook = default_walk_stack_root
+        root_iterator = get_root_iterator(gctransformer)
+        def walk_stack_root(callback, addr, end):
+            root_iterator.setcontext(NonConstant(llmemory.NULL))
+            gc = self.gc
+            while True:
+                addr = root_iterator.next(gc, addr, end)
+                if addr == llmemory.NULL:
+                    return
+                callback(gc, addr)
+                addr += sizeofaddr
+        self.rootstackhook = walk_stack_root
 
         self.shadow_stack_pool = ShadowStackPool(gcdata)
         rsd = gctransformer.root_stack_depth
@@ -331,6 +320,30 @@
                 raise MemoryError
 
 
+def get_root_iterator(gctransformer):
+    if hasattr(gctransformer, '_root_iterator'):
+        return gctransformer._root_iterator     # if already built
+    translator = gctransformer.translator
+    if (hasattr(translator, '_jit2gc') and
+            'root_iterator' in translator._jit2gc):
+        result = translator._jit2gc['root_iterator']
+    else:
+        class RootIterator(object):
+            def _freeze_(self):
+                return True
+            def setcontext(self, context):
+                pass
+            def next(self, gc, addr, end):
+                while addr != end:
+                    if gc.points_to_valid_gc_object(addr):
+                        return addr
+                    addr += sizeofaddr
+                return llmemory.NULL
+        result = RootIterator()
+    gctransformer._root_iterator = result
+    return result
+
+
 def get_shadowstackref(gctransformer):
     if hasattr(gctransformer, '_SHADOWSTACKREF'):
         return gctransformer._SHADOWSTACKREF
@@ -344,28 +357,17 @@
                                      rtti=True)
     SHADOWSTACKREFPTR.TO.become(SHADOWSTACKREF)
 
-    translator = gctransformer.translator
-    if hasattr(translator, '_jit2gc'):
-        gc = gctransformer.gcdata.gc
-        root_iterator = translator._jit2gc['root_iterator']
-        def customtrace(obj, prev):
-            obj = llmemory.cast_adr_to_ptr(obj, SHADOWSTACKREFPTR)
-            if not prev:
-                root_iterator.context = obj.context
-                next = obj.base
-            else:
-                next = prev + sizeofaddr
-            return root_iterator.next(gc, next, obj.top)
-    else:
-        def customtrace(obj, prev):
-            # a simple but not JIT-ready version
-            if not prev:
-                next = llmemory.cast_adr_to_ptr(obj, SHADOWSTACKREFPTR).base
-            else:
-                next = prev + sizeofaddr
-            if next == llmemory.cast_adr_to_ptr(obj, SHADOWSTACKREFPTR).top:
-                next = llmemory.NULL
-            return next
+    gc = gctransformer.gcdata.gc
+    root_iterator = get_root_iterator(gctransformer)
+
+    def customtrace(obj, prev):
+        obj = llmemory.cast_adr_to_ptr(obj, SHADOWSTACKREFPTR)
+        if not prev:
+            root_iterator.setcontext(obj.context)
+            next = obj.base
+        else:
+            next = prev + sizeofaddr
+        return root_iterator.next(gc, next, obj.top)
 
     CUSTOMTRACEFUNC = lltype.FuncType([llmemory.Address, llmemory.Address],
                                       llmemory.Address)