Add instance Lift ByteArray

2 files changed, 58 insertions, 1 deletions

diff --git a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
index 5acf96e011..08d2ea41bf 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
@@ -3,7 +3,7 @@
              RankNTypes, RoleAnnotations, ScopedTypeVariables,
              MagicHash, KindSignatures, PolyKinds, TypeApplications, DataKinds,
              GADTs, UnboxedTuples, UnboxedSums, TypeInType, TypeOperators,
-             Trustworthy, DeriveFunctor #-}
+             Trustworthy, DeriveFunctor, BangPatterns, RecordWildCards, ImplicitParams #-}
 
 {-# OPTIONS_GHC -fno-warn-inline-rule-shadowing #-}
 
@@ -61,11 +61,22 @@ import Prelude
 import Foreign.ForeignPtr
 import Foreign.C.String
 import Foreign.C.Types
+import GHC.Stack
 
 #if __GLASGOW_HASKELL__ >= 901
 import GHC.Types ( Levity(..) )
 #endif
 
+#if __GLASGOW_HASKELL__ >= 903
+import Data.Array.Byte (ByteArray(..))
+import GHC.Exts
+  ( ByteArray#, unsafeFreezeByteArray#, copyAddrToByteArray#, newByteArray#
+  , isByteArrayPinned#, isTrue#, sizeofByteArray#, unsafeCoerce#, byteArrayContents#
+  , copyByteArray#, newPinnedByteArray#)
+import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..))
+import GHC.ST (ST(..), runST)
+#endif
+
 -----------------------------------------------------
 --
 --              The Quasi class
@@ -1074,6 +1085,51 @@ instance Lift Addr# where
   lift x
     = return (LitE (StringPrimL (map (fromIntegral . ord) (unpackCString# x))))
 
+#if __GLASGOW_HASKELL__ >= 903
+
+-- |
+-- @since 2.19.0.0
+instance Lift ByteArray where
+  liftTyped x = unsafeCodeCoerce (lift x)
+  lift (ByteArray b) = return
+    (AppE (AppE (VarE addrToByteArrayName) (LitE (IntegerL (fromIntegral len))))
+      (LitE (BytesPrimL (Bytes ptr 0 (fromIntegral len)))))
+    where
+      len# = sizeofByteArray# b
+      len = I# len#
+      pb :: ByteArray#
+      !(ByteArray pb)
+        | isTrue# (isByteArrayPinned# b) = ByteArray b
+        | otherwise = runST $ ST $
+          \s -> case newPinnedByteArray# len# s of
+            (# s', mb #) -> case copyByteArray# b 0# mb 0# len# s' of
+              s'' -> case unsafeFreezeByteArray# mb s'' of
+                (# s''', ret #) -> (# s''', ByteArray ret #)
+      ptr :: ForeignPtr Word8
+      ptr = ForeignPtr (byteArrayContents# pb) (PlainPtr (unsafeCoerce# pb))
+
+
+-- We can't use a TH quote in this module because we're in the template-haskell
+-- package, so we conconct this quite defensive solution to make the correct name
+-- which will work if the package name or module name changes in future.
+addrToByteArrayName :: Name
+addrToByteArrayName = helper
+  where
+    helper :: HasCallStack => Name
+    helper =
+      case head (getCallStack ?callStack) of
+        (_, SrcLoc{..}) -> mkNameG_v srcLocPackage srcLocModule "addrToByteArray"
+
+
+addrToByteArray :: Int -> Addr# -> ByteArray
+addrToByteArray (I# len) addr = runST $ ST $
+  \s -> case newByteArray# len s of
+    (# s', mb #) -> case copyAddrToByteArray# addr mb 0# len s' of
+      s'' -> case unsafeFreezeByteArray# mb s'' of
+        (# s''', ret #) -> (# s''', ByteArray ret #)
+
+#endif
+
 instance Lift a => Lift (Maybe a) where
   liftTyped x = unsafeCodeCoerce (lift x)
 
diff --git a/libraries/template-haskell/changelog.md b/libraries/template-haskell/changelog.md
index 3965279f15..e3bcef509a 100644
--- a/libraries/template-haskell/changelog.md
+++ b/libraries/template-haskell/changelog.md
@@ -7,6 +7,7 @@
   * Add support for Overloaded Record Dot.
     Introduces `getFieldE :: Quote m => m Exp -> String -> m Exp` and
     `projectionE :: Quote m => [String] -> m Exp`.
+  * Add `instance Lift ByteArray`.
 
   * Add `PromotedInfixT` and `PromotedUInfixT`, which are analogs to `InfixT`
     and `UInfixT` that ensure that if a dynamically bound name (i.e. a name