From 44375faacf4c3ac007a03dfbae95016a6400b9e8 Mon Sep 17 00:00:00 2001
From: Matthew Craven <clyring@gmail.com>
Date: Tue, 30 Jan 2024 16:41:14 -0500
Subject: [PATCH] Use primops for unaligned writes when possible (#645)

* Use primops for unaligned writes when possible

* Check MIN_VERSION_base instead of __GLASGOW_HASKELL__
---
 Data/ByteString/Utils/UnalignedWrite.hs | 33 ++++++++++++++++++++++++-
 include/bytestring-cpp-macros.h         |  2 ++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/Data/ByteString/Utils/UnalignedWrite.hs b/Data/ByteString/Utils/UnalignedWrite.hs
index ff1c7393c..60daffbe0 100644
--- a/Data/ByteString/Utils/UnalignedWrite.hs
+++ b/Data/ByteString/Utils/UnalignedWrite.hs
@@ -1,5 +1,8 @@
 {-# LANGUAGE CPP #-}
 
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+
 #include "bytestring-cpp-macros.h"
 
 module Data.ByteString.Utils.UnalignedWrite
@@ -13,7 +16,34 @@ module Data.ByteString.Utils.UnalignedWrite
 import Foreign.Ptr
 import Data.Word
 
-#if HS_UNALIGNED_POKES_OK
+
+#if HS_UNALIGNED_ADDR_PRIMOPS_AVAILABLE
+import GHC.IO (IO(..))
+import GHC.Word (Word16(..), Word32(..), Word64(..))
+import GHC.Exts
+
+unalignedWriteU16 :: Word16 -> Ptr Word8 -> IO ()
+unalignedWriteU16 = coerce $ \(W16# x#) (Ptr p#) s
+  -> (# writeWord8OffAddrAsWord16# p# 0# x# s, () #)
+
+unalignedWriteU32 :: Word32 -> Ptr Word8 -> IO ()
+unalignedWriteU32 = coerce $ \(W32# x#) (Ptr p#) s
+  -> (# writeWord8OffAddrAsWord32# p# 0# x# s, () #)
+
+unalignedWriteU64 :: Word64 -> Ptr Word8 -> IO ()
+unalignedWriteU64 = coerce $ \(W64# x#) (Ptr p#) s
+  -> (# writeWord8OffAddrAsWord64# p# 0# x# s, () #)
+
+unalignedWriteFloat :: Float -> Ptr Word8 -> IO ()
+unalignedWriteFloat = coerce $ \(F# x#) (Ptr p#) s
+  -> (# writeWord8OffAddrAsFloat# p# 0# x# s, () #)
+
+unalignedWriteDouble :: Double -> Ptr Word8 -> IO ()
+unalignedWriteDouble = coerce $ \(D# x#) (Ptr p#) s
+  -> (# writeWord8OffAddrAsDouble# p# 0# x# s, () #)
+
+
+#elif HS_UNALIGNED_POKES_OK
 import Foreign.Storable
 
 unalignedWriteU16 :: Word16 -> Ptr Word8 -> IO ()
@@ -31,6 +61,7 @@ unalignedWriteFloat x p = poke (castPtr p) x
 unalignedWriteDouble :: Double -> Ptr Word8 -> IO ()
 unalignedWriteDouble x p = poke (castPtr p) x
 
+
 #else
 foreign import ccall unsafe "static fpstring.h fps_unaligned_write_u16"
   unalignedWriteU16 :: Word16 -> Ptr Word8 -> IO ()
diff --git a/include/bytestring-cpp-macros.h b/include/bytestring-cpp-macros.h
index e2008e5e2..aa04c168f 100644
--- a/include/bytestring-cpp-macros.h
+++ b/include/bytestring-cpp-macros.h
@@ -34,3 +34,5 @@ These operations were added in base-4.10.0, but due to
 https://gitlab.haskell.org/ghc/ghc/-/issues/16617 they
 are buggy with negative floats before ghc-8.10.
 */
+
+#define HS_UNALIGNED_ADDR_PRIMOPS_AVAILABLE MIN_VERSION_base(4,20,0)