From c6c8aff5b329f75145fcd9100faf9c8447c51ef9 Mon Sep 17 00:00:00 2001
From: David Johnson <djohnson.m@gmail.com>
Date: Tue, 5 Nov 2019 16:27:07 -0500
Subject: [PATCH 1/2] Enfore Complex on 'imag' and 'real' functions.

---
 src/ArrayFire/Arith.hs | 22 ++++++++++++----------
 src/ArrayFire/FFI.hs   | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/ArrayFire/Arith.hs b/src/ArrayFire/Arith.hs
index c1b4df4..e59b0b4 100644
--- a/src/ArrayFire/Arith.hs
+++ b/src/ArrayFire/Arith.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE TypeApplications    #-}
+{-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE ViewPatterns        #-}
 --------------------------------------------------------------------------------
@@ -28,6 +29,7 @@ import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegra
 
 import Data.Coerce
 import Data.Proxy
+import Data.Complex
 
 import ArrayFire.FFI
 import ArrayFire.Internal.Arith
@@ -1195,31 +1197,31 @@ cplx = flip op1 af_cplx
 
 -- | Execute real
 --
--- >>> A.real (A.vector @Double 10 [1..])
+-- >>> A.real (A.scalar @(Complex Double) (10 :+ 11)) :: Array Double
 -- ArrayFire Array
 -- [10 1 1 1]
---    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
+--    10.0000
 real
-  :: AFType a
-  => Array a
+  :: (AFType a, AFType (Complex b))
+  => Array (Complex b)
   -- ^ Input array
   -> Array a
   -- ^ Result of calling 'real'
-real = flip op1 af_real
+real = flip op1d af_real
 
 -- | Execute imag
 --
--- >>> A.imag (A.vector @Double 10 [1..])
+-- >>> A.imag (A.scalar @(Complex Double) (10 :+ 11)) :: Array Double
 -- ArrayFire Array
 -- [10 1 1 1]
---    0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000
+--    11.0000
 imag
-  :: AFType a
-  => Array a
+  :: (AFType a, AFType (Complex b))
+  => Array (Complex b)
   -- ^ Input array
   -> Array a
   -- ^ Result of calling 'imag'
-imag = flip op1 af_imag
+imag = flip op1d af_imag
 
 -- | Execute conjg
 --
diff --git a/src/ArrayFire/FFI.hs b/src/ArrayFire/FFI.hs
index df3eb2f..4e4609d 100644
--- a/src/ArrayFire/FFI.hs
+++ b/src/ArrayFire/FFI.hs
@@ -220,6 +220,22 @@ opw1 (Window fptr) op
          throwAFError =<< op p ptr
          peek p
 
+op1d
+  :: Array a
+  -> (Ptr AFArray -> AFArray -> IO AFErr)
+  -> Array b
+{-# NOINLINE op1d #-}
+op1d (Array fptr1) op =
+  unsafePerformIO $ do
+    withForeignPtr fptr1 $ \ptr1 -> do
+      ptr <-
+        alloca $ \ptrInput -> do
+          throwAFError =<< op ptrInput ptr1
+          peek ptrInput
+      fptr <- newForeignPtr af_release_array_finalizer ptr
+      pure (Array fptr)
+
+
 op1
   :: Array a
   -> (Ptr AFArray -> AFArray -> IO AFErr)

From a35753a72385d6d8445126b6c490f8d2f1639316 Mon Sep 17 00:00:00 2001
From: David Johnson <djohnson.m@gmail.com>
Date: Tue, 5 Nov 2019 21:43:56 -0500
Subject: [PATCH 2/2] Constrain real/imag on RealFrac.

---
 src/ArrayFire/Arith.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/ArrayFire/Arith.hs b/src/ArrayFire/Arith.hs
index e59b0b4..fbe2e0c 100644
--- a/src/ArrayFire/Arith.hs
+++ b/src/ArrayFire/Arith.hs
@@ -25,7 +25,7 @@
 --------------------------------------------------------------------------------
 module ArrayFire.Arith where
 
-import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegral, Real)
+import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegral, Real, RealFrac)
 
 import Data.Coerce
 import Data.Proxy
@@ -1202,7 +1202,7 @@ cplx = flip op1 af_cplx
 -- [10 1 1 1]
 --    10.0000
 real
-  :: (AFType a, AFType (Complex b))
+  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
   => Array (Complex b)
   -- ^ Input array
   -> Array a
@@ -1216,7 +1216,7 @@ real = flip op1d af_real
 -- [10 1 1 1]
 --    11.0000
 imag
-  :: (AFType a, AFType (Complex b))
+  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
   => Array (Complex b)
   -- ^ Input array
   -> Array a