Add instance Lift ByteArray

5 files changed, 72 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
diff --git a/testsuite/tests/th/Lift_ByteArray.hs b/testsuite/tests/th/Lift_ByteArray.hs
new file mode 100644
index 0000000000..6463e972e5
--- /dev/null
+++ b/testsuite/tests/th/Lift_ByteArray.hs
@@ -0,0 +1,12 @@
+{-# LANGUAGE OverloadedLists #-}
+{-# LANGUAGE TemplateHaskell #-}
+module Main where
+
+import Data.Array.Byte
+import Language.Haskell.TH.Syntax
+
+
+main = print $(
+    let liftBA :: ByteArray -> Q Exp
+        liftBA = lift
+    in liftBA [0..128])
diff --git a/testsuite/tests/th/Lift_ByteArray.stdout b/testsuite/tests/th/Lift_ByteArray.stdout
new file mode 100644
index 0000000000..b504ccbe90
--- /dev/null
+++ b/testsuite/tests/th/Lift_ByteArray.stdout
@@ -0,0 +1 @@
+[0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4a, 0x4b, 0x4c, 0x4d, 0x4e, 0x4f, 0x50, 0x51, 0x52, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5a, 0x5b, 0x5c, 0x5d, 0x5e, 0x5f, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6a, 0x6b, 0x6c, 0x6d, 0x6e, 0x6f, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7a, 0x7b, 0x7c, 0x7d, 0x7e, 0x7f, 0x80]
diff --git a/testsuite/tests/th/all.T b/testsuite/tests/th/all.T
index 97e5700607..276faf84d8 100644
--- a/testsuite/tests/th/all.T
+++ b/testsuite/tests/th/all.T
@@ -550,3 +550,4 @@ test('T15433a', [extra_files(['T15433_aux.hs'])], multimod_compile_fail, ['T1543
 test('T15433b', [extra_files(['T15433_aux.hs'])], multimod_compile, ['T15433b', '-v0'])
 test('T20711', normal, compile_and_run, [''])
 test('T20868', normal, compile_and_run, [''])
+test('Lift_ByteArray', normal, compile_and_run, [''])