[pypy-commit] pypy default: Support 'llhelper(F, multiple_functions)', as long as they all have the

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r85118:6709ea850047
Date: 2016-06-13 10:18 +0200
http://bitbucket.org/pypy/pypy/changeset/6709ea850047/

Log:	Support 'llhelper(F, multiple_functions)', as long as they all have
	the right signature.

diff --git a/rpython/rtyper/annlowlevel.py b/rpython/rtyper/annlowlevel.py
--- a/rpython/rtyper/annlowlevel.py
+++ b/rpython/rtyper/annlowlevel.py
@@ -348,19 +348,30 @@
     _about_ = llhelper
 
     def compute_result_annotation(self, s_F, s_callable):
+        from rpython.annotator.description import FunctionDesc
         assert s_F.is_constant()
-        assert s_callable.is_constant()
+        assert isinstance(s_callable, annmodel.SomePBC)
         F = s_F.const
         FUNC = F.TO
         args_s = [lltype_to_annotation(T) for T in FUNC.ARGS]
-        key = (llhelper, s_callable.const)
-        s_res = self.bookkeeper.emulate_pbc_call(key, s_callable, args_s)
-        assert lltype_to_annotation(FUNC.RESULT).contains(s_res)
+        for desc in s_callable.descriptions:
+            assert isinstance(desc, FunctionDesc)
+            assert desc.pyobj is not None
+            if s_callable.is_constant():
+                assert s_callable.const is desc.pyobj
+            key = (llhelper, desc.pyobj)
+            s_res = self.bookkeeper.emulate_pbc_call(key, s_callable, args_s)
+            assert lltype_to_annotation(FUNC.RESULT).contains(s_res)
         return SomePtr(F)
 
     def specialize_call(self, hop):
         hop.exception_cannot_occur()
-        return hop.args_r[1].get_unique_llfn()
+        if hop.args_s[1].is_constant():
+            return hop.args_r[1].get_unique_llfn()
+        else:
+            F = hop.args_s[0].const
+            assert hop.args_r[1].lowleveltype == F
+            return hop.inputarg(hop.args_r[1], 1)
 
 # ____________________________________________________________
 
diff --git a/rpython/rtyper/test/test_llann.py b/rpython/rtyper/test/test_llann.py
--- a/rpython/rtyper/test/test_llann.py
+++ b/rpython/rtyper/test/test_llann.py
@@ -462,6 +462,30 @@
     res = interpret(h, [8, 5, 2])
     assert res == 99
 
+def test_llhelper_multiple_functions():
+    S = GcStruct('S', ('x', Signed), ('y', Signed))
+    def f(s):
+        return s.x - s.y
+    def g(s):
+        return s.x + s.y
+
+    F = Ptr(FuncType([Ptr(S)], Signed))
+
+    myfuncs = [f, g]
+
+    def h(x, y, z):
+        s = malloc(S)
+        s.x = x
+        s.y = y
+        fptr = llhelper(F, myfuncs[z])
+        assert typeOf(fptr) == F
+        return fptr(s)
+
+    res = interpret(h, [80, 5, 0])
+    assert res == 75
+    res = interpret(h, [80, 5, 1])
+    assert res == 85
+
 
 def test_cast_instance_to_base_ptr():
     class A: