Call support for functions with complex args/return types

4 files changed, 77 insertions, 7 deletions

diff --git a/pypy/module/_cffi_backend/realize_c_type.py b/pypy/module/_cffi_backend/realize_c_type.py
index a8b4ec62c6..3afe7ced7b 100644
--- a/pypy/module/_cffi_backend/realize_c_type.py
+++ b/pypy/module/_cffi_backend/realize_c_type.py
@@ -8,6 +8,7 @@ from pypy.interpreter.baseobjspace import W_Root
 from pypy.module import _cffi_backend
 from pypy.module._cffi_backend.ctypeobj import W_CType
 from pypy.module._cffi_backend import cffi_opcode, newtype, ctypestruct
+from pypy.module._cffi_backend import ctypeprim
 from pypy.module._cffi_backend import parse_c_type
 
 
@@ -70,6 +71,8 @@ class RealizeCache:
         "uint_fast64_t",
         "intmax_t",
         "uintmax_t",
+        "float _Complex",
+        "double _Complex",
         ]
     assert len(NAMES) == cffi_opcode._NUM_PRIM
 
@@ -209,7 +212,7 @@ class W_RawFuncType(W_Root):
         # which the struct args are replaced with ptr-to- struct, and
         # a struct return value is replaced with a hidden first arg of
         # type ptr-to-struct.  This is how recompiler.py produces
-        # trampoline functions for PyPy.
+        # trampoline functions for PyPy.  (Same with complex numbers.)
         if self.nostruct_ctype is None:
             fargs, fret, ellipsis, abi = self._unpack(ffi)
             # 'locs' will be a string of the same length as the final fargs,
@@ -218,11 +221,13 @@ class W_RawFuncType(W_Root):
             locs = ['\x00'] * len(fargs)
             for i in range(len(fargs)):
                 farg = fargs[i]
-                if isinstance(farg, ctypestruct.W_CTypeStructOrUnion):
+                if (isinstance(farg, ctypestruct.W_CTypeStructOrUnion) or
+                    isinstance(farg, ctypeprim.W_CTypePrimitiveComplex)):
                     farg = newtype.new_pointer_type(ffi.space, farg)
                     fargs[i] = farg
                     locs[i] = 'A'
-            if isinstance(fret, ctypestruct.W_CTypeStructOrUnion):
+            if (isinstance(fret, ctypestruct.W_CTypeStructOrUnion) or
+                isinstance(fret, ctypeprim.W_CTypePrimitiveComplex)):
                 fret = newtype.new_pointer_type(ffi.space, fret)
                 fargs = [fret] + fargs
                 locs = ['R'] + locs
diff --git a/pypy/module/_cffi_backend/test/_backend_test_c.py b/pypy/module/_cffi_backend/test/_backend_test_c.py
index 47ba162186..575f0d5741 100644
--- a/pypy/module/_cffi_backend/test/_backend_test_c.py
+++ b/pypy/module/_cffi_backend/test/_backend_test_c.py
@@ -1,7 +1,7 @@
 # ____________________________________________________________
 
 import sys
-assert __version__ == "1.10.0", ("This test_c.py file is for testing a version"
+assert __version__ == "1.11.0", ("This test_c.py file is for testing a version"
                                  " of cffi that differs from the one that we"
                                  " get from 'import _cffi_backend'")
 if sys.version_info < (3,):
diff --git a/pypy/module/_cffi_backend/test/test_recompiler.py b/pypy/module/_cffi_backend/test/test_recompiler.py
index 7485a4a925..bf72765392 100644
--- a/pypy/module/_cffi_backend/test/test_recompiler.py
+++ b/pypy/module/_cffi_backend/test/test_recompiler.py
@@ -1819,6 +1819,68 @@ class AppTestRecompiler:
         assert lib.f.__get__(42) is lib.f
         assert lib.f.__get__(42, int) is lib.f
 
+    def test_function_returns_float_complex(self):
+        import sys
+        if sys.platform == 'win32':
+            skip("MSVC may not support _Complex")
+        ffi, lib = self.prepare(
+            "float _Complex f1(float a, float b);",
+            "test_function_returns_float_complex", """
+            #include <complex.h>
+            static float _Complex f1(float a, float b) { return a + I*2.0*b; }
+        """, min_version=(1, 11, 0))
+        result = lib.f1(1.25, 5.1)
+        assert type(result) == complex
+        assert result.real == 1.25   # exact
+        assert (result.imag != 2*5.1) and (abs(result.imag - 2*5.1) < 1e-5) # inexact
+
+    def test_function_returns_double_complex(self):
+        import sys
+        if sys.platform == 'win32':
+            skip("MSVC may not support _Complex")
+        ffi, lib = self.prepare(
+            "double _Complex f1(double a, double b);",
+            "test_function_returns_double_complex", """
+            #include <complex.h>
+            static double _Complex f1(double a, double b) { return a + I*2.0*b; }
+        """, min_version=(1, 11, 0))
+        result = lib.f1(1.25, 5.1)
+        assert type(result) == complex
+        assert result.real == 1.25   # exact
+        assert result.imag == 2*5.1  # exact
+
+    def test_function_argument_float_complex(self):
+        import sys
+        if sys.platform == 'win32':
+            skip("MSVC may not support _Complex")
+        ffi, lib = self.prepare(
+            "float f1(float _Complex x);",
+            "test_function_argument_float_complex", """
+            #include <complex.h>
+            static float f1(float _Complex x) { return cabsf(x); }
+        """, min_version=(1, 11, 0))
+        x = complex(12.34, 56.78)
+        result = lib.f1(x)
+        assert abs(result - abs(x)) < 1e-5
+        result2 = lib.f1(ffi.cast("float _Complex", x))
+        assert result2 == result
+
+    def test_function_argument_double_complex(self):
+        import sys
+        if sys.platform == 'win32':
+            skip("MSVC may not support _Complex")
+        ffi, lib = self.prepare(
+            "double f1(double _Complex);",
+            "test_function_argument_double_complex", """
+            #include <complex.h>
+            static double f1(double _Complex x) { return cabs(x); }
+        """, min_version=(1, 11, 0))
+        x = complex(12.34, 56.78)
+        result = lib.f1(x)
+        assert abs(result - abs(x)) < 1e-11
+        result2 = lib.f1(ffi.cast("double _Complex", x))
+        assert result2 == result
+
     def test_typedef_array_dotdotdot(self):
         ffi, lib = self.prepare("""
             typedef int foo_t[...], bar_t[...];
diff --git a/pypy/module/_cffi_backend/wrapper.py b/pypy/module/_cffi_backend/wrapper.py
index ea3a468c13..7967c451e2 100644
--- a/pypy/module/_cffi_backend/wrapper.py
+++ b/pypy/module/_cffi_backend/wrapper.py
@@ -8,6 +8,7 @@ from rpython.rlib import jit
 from pypy.module._cffi_backend.cdataobj import W_CData
 from pypy.module._cffi_backend.cdataobj import W_CDataPtrToStructOrUnion
 from pypy.module._cffi_backend.ctypeptr import W_CTypePtrOrArray
+from pypy.module._cffi_backend.ctypeptr import W_CTypePointer
 from pypy.module._cffi_backend.ctypefunc import W_CTypeFunc
 from pypy.module._cffi_backend.ctypestruct import W_CTypeStructOrUnion
 from pypy.module._cffi_backend import allocator
@@ -83,8 +84,9 @@ class W_FunctionWrapper(W_Root):
                 #
                 ctype._call(self.fnptr, args_w)    # returns w_None
                 #
-                assert isinstance(w_result_cdata, W_CDataPtrToStructOrUnion)
-                return w_result_cdata.structobj
+                ctyperesptr = w_result_cdata.ctype
+                assert isinstance(ctyperesptr, W_CTypePointer)
+                return w_result_cdata._do_getitem(ctyperesptr, 0)
             else:
                 args_w = args_w[:]
                 prepare_args(space, rawfunctype, args_w, 0)
@@ -109,13 +111,14 @@ class W_FunctionWrapper(W_Root):
 @jit.unroll_safe
 def prepare_args(space, rawfunctype, args_w, start_index):
     # replaces struct/union arguments with ptr-to-struct/union arguments
+    # as well as complex numbers
     locs = rawfunctype.nostruct_locs
     fargs = rawfunctype.nostruct_ctype.fargs
     for i in range(start_index, len(locs)):
         if locs[i] != 'A':
             continue
         w_arg = args_w[i]
-        farg = fargs[i]      # <ptr to struct/union>
+        farg = fargs[i]      # <ptr to struct/union/complex>
         assert isinstance(farg, W_CTypePtrOrArray)
         if isinstance(w_arg, W_CData) and w_arg.ctype is farg.ctitem:
             # fast way: we are given a W_CData "struct", so just make