[pypy-svn] r67720 - in pypy/trunk/pypy/rlib: . test

[antocuni at codespeak.net]

Author: antocuni
Date: Wed Sep 16 17:32:15 2009
New Revision: 67720

Modified:
   pypy/trunk/pypy/rlib/jit.py
   pypy/trunk/pypy/rlib/test/test_jit.py
Log:
make sure that hooks attached to the jitdriver are annotated at the right
time. So far they were annotated only when the codewriter tried to get them,
i.e. when the main annotation/rtyping phase was already over.  On lltype it
worked by chance, and on ootype it caused problems.



Modified: pypy/trunk/pypy/rlib/jit.py
==============================================================================
--- pypy/trunk/pypy/rlib/jit.py	(original)
+++ pypy/trunk/pypy/rlib/jit.py	Wed Sep 16 17:32:15 2009
@@ -203,8 +203,28 @@
         for name in driver.greens:
             s_green_key = kwds_s['s_' + name]
             s_green_key.hash()      # force the hash cache to appear
+
+        self.annotate_hooks(**kwds_s)
         return annmodel.s_None
 
+    def annotate_hooks(self, **kwds_s):
+        driver = self.instance.im_self
+        self.annotate_hook(driver.can_inline, driver.greens, **kwds_s)
+        self.annotate_hook(driver.get_printable_location, driver.greens, **kwds_s)
+        self.annotate_hook(driver.leave, driver.greens + driver.reds, **kwds_s)
+
+    def annotate_hook(self, func, variables, **kwds_s):
+        if func is None:
+            return
+        bk = self.bookkeeper
+        s_func = bk.immutablevalue(func)
+        uniquekey = 'jitdriver.%s' % func.func_name
+        args_s = []
+        for name in variables:
+            s_arg = kwds_s['s_' + name]
+            args_s.append(s_arg)
+        bk.emulate_pbc_call(uniquekey, s_func, args_s)
+
     def specialize_call(self, hop, **kwds_i):
         # XXX to be complete, this could also check that the concretetype
         # of the variables are the same for each of the calls.

Modified: pypy/trunk/pypy/rlib/test/test_jit.py
==============================================================================
--- pypy/trunk/pypy/rlib/test/test_jit.py	(original)
+++ pypy/trunk/pypy/rlib/test/test_jit.py	Wed Sep 16 17:32:15 2009
@@ -1,9 +1,10 @@
 import py
-from pypy.rlib.jit import hint, we_are_jitted
+from pypy.rlib.jit import hint, we_are_jitted, JitDriver
 from pypy.translator.translator import TranslationContext, graphof
 from pypy.rpython.test.tool import BaseRtypingTest, LLRtypeMixin, OORtypeMixin
+from pypy.rpython.lltypesystem import lltype
 
-class TestJIT(BaseRtypingTest, LLRtypeMixin):
+class BaseTestJIT(BaseRtypingTest):
     def test_hint(self):
         def f():
             x = hint(5, hello="world")
@@ -22,3 +23,43 @@
         res = self.interpret(f, [4])
         assert res == 5
 
+
+    def test_annotate_hooks(self):
+        
+        def can_inline(m): pass
+        def get_printable_location(m): pass
+        def leave(m, n): pass
+        
+        myjitdriver = JitDriver(greens=['m'], reds=['n'],
+                                can_inline=can_inline,
+                                get_printable_location=get_printable_location,
+                                leave=leave)
+        def fn(n):
+            m = 42.5
+            while n > 0:
+                myjitdriver.can_enter_jit(m=m, n=n)
+                myjitdriver.jit_merge_point(m=m, n=n)
+                n -= 1
+            return n
+
+        t, rtyper, fngraph = self.gengraph(fn, [int])
+
+        def getargs(func):
+            for graph in t.graphs:
+                if getattr(graph, 'func', None) is func:
+                    return [v.concretetype for v in graph.getargs()]
+            raise Exception, 'function %r has not been annotated' % func
+
+        leave_args = getargs(leave)
+        assert leave_args == [lltype.Float, lltype.Signed]
+
+        can_inline_args = getargs(can_inline)
+        get_printable_location_args = getargs(get_printable_location)
+        assert can_inline_args == get_printable_location_args == [lltype.Float]
+
+
+class TestJITLLtype(BaseTestJIT, LLRtypeMixin):
+    pass
+
+class TestJITOOtype(BaseTestJIT, OORtypeMixin):
+    pass