[pypy-svn] r64943 - in pypy/branch/pyjitpl5/pypy/translator: . test

[arigo at codespeak.net]

Author: arigo
Date: Fri May  1 14:56:21 2009
New Revision: 64943

Modified:
   pypy/branch/pyjitpl5/pypy/translator/test/test_unsimplify.py
   pypy/branch/pyjitpl5/pypy/translator/unsimplify.py
Log:
Add a helper, unsimplify.call_final_function().


Modified: pypy/branch/pyjitpl5/pypy/translator/test/test_unsimplify.py
==============================================================================
--- pypy/branch/pyjitpl5/pypy/translator/test/test_unsimplify.py	(original)
+++ pypy/branch/pyjitpl5/pypy/translator/test/test_unsimplify.py	Fri May  1 14:56:21 2009
@@ -1,12 +1,15 @@
+import os
 from pypy.translator.translator import TranslationContext, graphof
-from pypy.translator.unsimplify import split_block
+from pypy.translator.unsimplify import split_block, call_final_function
 from pypy.rpython.llinterp import LLInterpreter
 from pypy.objspace.flow.model import checkgraph
+from pypy.rlib.objectmodel import we_are_translated
+from pypy.tool.udir import udir
 
-def translate(func, argtypes):
+def translate(func, argtypes, type_system="lltype"):
     t = TranslationContext()
     t.buildannotator().build_types(func, argtypes)
-    t.buildrtyper().specialize()
+    t.buildrtyper(type_system=type_system).specialize()
     return graphof(t, func), t
 
 def test_split_blocks_simple():
@@ -66,3 +69,21 @@
         result = interp.eval_graph(graph, [2])
         assert result == 2
     
+def test_call_final_function():
+    tmpfile = str(udir.join('test_call_final_function'))
+    for type_system in ['lltype', 'ootype']:
+        def f(x):
+            return x * 6
+        def goodbye_world():
+            if we_are_translated():
+                fd = os.open(tmpfile, os.O_WRONLY | os.O_CREAT, 0)
+                os.close(fd)
+        graph, t = translate(f, [int], type_system)
+        call_final_function(t, goodbye_world)
+        #
+        if os.path.exists(tmpfile):
+            os.unlink(tmpfile)
+        interp = LLInterpreter(t.rtyper)
+        result = interp.eval_graph(graph, [7])
+        assert result == 42
+        assert os.path.isfile(tmpfile)

Modified: pypy/branch/pyjitpl5/pypy/translator/unsimplify.py
==============================================================================
--- pypy/branch/pyjitpl5/pypy/translator/unsimplify.py	(original)
+++ pypy/branch/pyjitpl5/pypy/translator/unsimplify.py	Fri May  1 14:56:21 2009
@@ -154,3 +154,27 @@
                 break
     if links_to_start_block:
         insert_empty_startblock(None, graph)
+
+def call_final_function(translator, final_func):
+    """When the program finishes normally, call 'final_func()'."""
+    from pypy.annotation import model as annmodel
+    from pypy.rpython.lltypesystem import lltype
+    from pypy.rpython.annlowlevel import MixLevelHelperAnnotator
+
+    annhelper = MixLevelHelperAnnotator(translator.rtyper)
+    c_final_func = annhelper.constfunc(final_func, [], annmodel.s_None)
+    annhelper.finish()
+
+    entry_point = translator.graphs[0]
+    v = copyvar(translator.annotator, entry_point.getreturnvar())
+    extrablock = Block([v])
+    v_none = varoftype(lltype.Void)
+    newop = SpaceOperation('direct_call', [c_final_func], v_none)
+    extrablock.operations = [newop]
+    extrablock.closeblock(Link([v], entry_point.returnblock))
+    for block in entry_point.iterblocks():
+        if block is not extrablock:
+            for link in block.exits:
+                if link.target is entry_point.returnblock:
+                    link.target = extrablock
+    checkgraph(entry_point)