From e13d94f7af25a7c4c4c12d7d7524bc55beb58e40 Mon Sep 17 00:00:00 2001
From: Travis Whitaker <pi.boy.travis@gmail.com>
Date: Fri, 5 Apr 2024 21:47:22 -0700
Subject: [PATCH] Add Data.Complex instances

---
 src/Data/Aeson/Types/FromJSON.hs |  9 +++++++++
 src/Data/Aeson/Types/ToJSON.hs   | 11 +++++++++++
 tests/PropertyRoundTrip.hs       |  3 +++
 3 files changed, 23 insertions(+)

diff --git a/src/Data/Aeson/Types/FromJSON.hs b/src/Data/Aeson/Types/FromJSON.hs
index 3f766e43a..0f90efd83 100644
--- a/src/Data/Aeson/Types/FromJSON.hs
+++ b/src/Data/Aeson/Types/FromJSON.hs
@@ -93,6 +93,7 @@ import Data.Aeson.Types.Internal
 import Data.Aeson.Decoding.ByteString.Lazy
 import Data.Aeson.Decoding.Conversion (unResult, toResultValue, lbsSpace)
 import Data.Bits (unsafeShiftR)
+import Data.Complex (Complex(..))
 import Data.Fixed (Fixed, HasResolution (resolution), Nano)
 import Data.Functor.Compose (Compose(..))
 import Data.Functor.Identity (Identity(..))
@@ -1720,6 +1721,14 @@ instance (FromJSON a, Integral a) => FromJSON (Ratio a) where
             then fail "Ratio denominator was 0"
             else pure $ numerator % denominator
 
+instance FromJSON a => FromJSON (Complex a) where
+    parseJSON = withArray "Complex" $ \c ->
+        let n = V.length c
+        in if n == 2
+           then (:+) <$> parseJSONElemAtIndex parseJSON 0 c
+                     <*> parseJSONElemAtIndex parseJSON 1 c
+           else fail $ "cannot unpack array of length " ++ show n ++ "into a Complex"
+
 -- | This instance includes a bounds check to prevent maliciously
 -- large inputs to fill up the memory of the target system. You can
 -- newtype 'Scientific' and provide your own instance using
diff --git a/src/Data/Aeson/Types/ToJSON.hs b/src/Data/Aeson/Types/ToJSON.hs
index 7de887593..ce071c324 100644
--- a/src/Data/Aeson/Types/ToJSON.hs
+++ b/src/Data/Aeson/Types/ToJSON.hs
@@ -67,6 +67,7 @@ import Data.Aeson.Types.Internal
 import qualified Data.Aeson.Key as Key
 import qualified Data.Aeson.KeyMap as KM
 import Data.Bits (unsafeShiftR)
+import Data.Complex (Complex(..))
 import Data.DList (DList)
 import Data.Fixed (Fixed, HasResolution, Nano)
 import Data.Foldable (toList)
@@ -1421,6 +1422,16 @@ instance (ToJSON a, Integral a) => ToJSON (Ratio a) where
         "numerator" .= numerator r <>
         "denominator" .= denominator r
 
+instance ToJSON a => ToJSON (Complex a) where
+    toJSON (i :+ q) = Array $ V.create $ do
+        mv <- VM.unsafeNew 2
+        VM.unsafeWrite mv 0 (toJSON i)
+        VM.unsafeWrite mv 1 (toJSON q)
+        return mv
+    toEncoding (i :+ q) = E.list id
+        [ toEncoding i
+        , toEncoding q
+        ]
 
 instance HasResolution a => ToJSON (Fixed a) where
     toJSON = Number . realToFrac
diff --git a/tests/PropertyRoundTrip.hs b/tests/PropertyRoundTrip.hs
index f21a3da86..5110a2e2f 100644
--- a/tests/PropertyRoundTrip.hs
+++ b/tests/PropertyRoundTrip.hs
@@ -7,6 +7,7 @@ import Prelude.Compat
 
 import Control.Applicative (Const)
 import Data.Aeson.Types
+import Data.Complex (Complex(..))
 import Data.DList (DList)
 import Data.List.NonEmpty (NonEmpty)
 import Data.Map (Map)
@@ -80,6 +81,8 @@ roundTripTests =
     , testProperty "Seq" $ roundTripEq @(Seq Int)
     , testProperty "Rational" $ roundTripEq @Rational
     , testProperty "Ratio Int" $ roundTripEq @(Ratio Int)
+    , testProperty "Complex Float" $ roundTripEq @(Complex Float)
+    , testProperty "Complex Double" $ roundTripEq @(Complex Double)
     , testProperty "UUID" $ roundTripEq @UUID.UUID
     , testProperty "These" $ roundTripEq @(These Char Bool)
     , testProperty "Fix" $ roundTripEq @(F.Fix (These Char))