Merge pull request #906 from haskell/autolet

Use let-insertion to pre-encode TH-generated keys.

Author: phadej

aeson.cabal

@@ -87,6 +87,7 @@ library
     Data.Aeson.Internal.ByteString
     Data.Aeson.Internal.Functions
     Data.Aeson.Internal.Text
+    Data.Aeson.Internal.TH
     Data.Aeson.Parser.Time
     Data.Aeson.Parser.Unescape
     Data.Aeson.Types.Class

benchmarks/aeson-benchmarks.cabal

@@ -75,6 +75,7 @@ library
     Data.Aeson.Internal.ByteString
     Data.Aeson.Internal.Functions
     Data.Aeson.Internal.Text
+    Data.Aeson.Internal.TH
     Data.Aeson.Internal.Time
     Data.Aeson.Key
     Data.Aeson.KeyMap

benchmarks/bench/Auto/T/D.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE TemplateHaskell #-}
+-- {-# OPTIONS_GHC -ddump-splices #-}
 
 module Auto.T.D where
 

src/Data/Aeson/Encoding/Internal.hs

@@ -16,6 +16,7 @@ module Data.Aeson.Encoding.Internal
     , pairs
     , pair
     , pairStr
+    , unsafePairSBS
     , pair'
     -- * Predicates
     , nullEncoding
@@ -65,6 +66,7 @@ import Prelude.Compat
 
 import Data.Aeson.Types.Internal (Value, Key)
 import Data.ByteString.Builder (Builder, char7, toLazyByteString)
+import Data.ByteString.Short (ShortByteString)
 import qualified Data.Aeson.Key as Key
 import Data.Int
 import Data.Scientific (Scientific)
@@ -142,6 +144,19 @@ pairStr name val = pair' (string name) val
 pair' :: Encoding' Key -> Encoding -> Series
 pair' name val = Value $ retagEncoding $ retagEncoding name >< colon >< val
 
+-- | A variant of a 'pair' where key is already encoded
+-- including the quotes and colon.
+--
+-- @
+-- 'pair' "foo" v = 'unsafePair' "\\"foo\\":" v
+-- @
+--
+-- @since 2.0.3.0
+--
+unsafePairSBS :: ShortByteString -> Encoding -> Series
+unsafePairSBS k v = Value $ retagEncoding $ Encoding (B.shortByteString k) >< v
+{-# INLINE unsafePairSBS #-}
+
 instance Semigroup Series where
     Empty   <> a       = a
     a       <> Empty   = a

src/Data/Aeson/Internal/ByteString.hs

@@ -1,14 +1,24 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE MagicHash #-}
+{-# LANGUAGE TemplateHaskellQuotes #-}
 module Data.Aeson.Internal.ByteString (
     mkBS, 
     withBS,
+    liftSBS,
 ) where
 
 import Data.ByteString.Internal (ByteString (..))
 import Data.Word (Word8)
 import Foreign.ForeignPtr (ForeignPtr)
+import Data.ByteString.Short (ShortByteString, fromShort)
+import GHC.Exts (Addr#, Ptr (Ptr))
+import Data.ByteString.Internal (accursedUnutterablePerformIO)
+import Data.ByteString.Short.Internal (createFromPtr)
+
+import qualified Data.ByteString as BS
+import qualified Language.Haskell.TH.Lib as TH
+import qualified Language.Haskell.TH.Syntax as TH
 
 #if !MIN_VERSION_bytestring(0,11,0)
 #if MIN_VERSION_base(4,10,0)
@@ -56,3 +66,22 @@ plusForeignPtr (ForeignPtr addr guts) (I# offset) = ForeignPtr (plusAddr# addr o
  #-}
 #endif
 #endif
+
+liftSBS :: ShortByteString -> TH.ExpQ
+#if MIN_VERSION_template_haskell(2,16,0)
+liftSBS sbs = withBS bs $ \ptr len -> [| unsafePackLenLiteral |]
+    `TH.appE` TH.litE (TH.integerL (fromIntegral len))
+    `TH.appE` TH.litE (TH.BytesPrimL $ TH.Bytes ptr 0 (fromIntegral len))
+    where
+      bs = fromShort sbs
+#else
+liftSBS sbs = withBS bs $ \_ len -> [| unsafePackLenLiteral |]
+    `TH.appE` TH.litE (TH.integerL (fromIntegral len))
+    `TH.appE` TH.litE (TH.StringPrimL $ BS.unpack bs)
+    where
+      bs = fromShort sbs
+#endif
+
+unsafePackLenLiteral :: Int -> Addr# -> ShortByteString
+unsafePackLenLiteral len addr# =
+    accursedUnutterablePerformIO $ createFromPtr (Ptr addr#) len

src/Data/Aeson/Internal/TH.hs

@@ -0,0 +1,71 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE RankNTypes #-}
+module Data.Aeson.Internal.TH (
+    letrecE,
+    autoletE,
+) where
+
+import Data.IORef              (IORef, atomicModifyIORef, newIORef, readIORef)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Language.Haskell.TH     (varE, ExpQ, Name, Q, newName, runIO)
+import System.IO.Unsafe        (unsafeInterleaveIO)
+
+import qualified Data.Map as Map
+import qualified Language.Haskell.TH.Syntax as TH
+
+letrecE
+    :: forall a. Ord a
+    => ((a -> ExpQ) -> (a -> ExpQ))
+    -> ((a -> ExpQ) -> ExpQ)
+    -> ExpQ
+letrecE f g = do
+    ref <- runIO $ newIORef Map.empty
+    expr <- g (loop ref)
+    bindings <- runIO $ readIORef ref
+    mkLet bindings expr
+  where
+    mkLet :: Map.Map a (Name, TH.Exp) -> TH.Exp -> ExpQ
+    mkLet bindings expr = do
+        return $ TH.LetE
+            [ TH.ValD (TH.VarP name) (TH.NormalB code) []
+            | (_, (name, code)) <- Map.toList bindings
+            ]
+            expr
+
+    loop :: IORef (Map.Map a (Name, TH.Exp)) -> a -> ExpQ
+    loop ref y = do
+        memo <- runIO $ readIORef ref
+        case Map.lookup y memo of
+            Nothing -> do
+                name <- newName $ "_let" ++ show (Map.size memo)
+                _ <- mfix_ $ \yCode -> do
+                    runIO $ atomicModifyIORef ref $ \m -> (Map.insert y (name, yCode) m, ())
+                    f (loop ref) y
+                varE name
+
+            Just (name, _) ->
+                varE name
+
+-- | Better 'letE'.
+autoletE
+    :: Ord a
+    => (a -> ExpQ)            -- ^ what bindings are
+    -> ((a -> ExpQ) -> ExpQ)  -- ^ expression with a function to generate bindings
+    -> ExpQ
+autoletE f = letrecE (const f)
+
+-------------------------------------------------------------------------------
+-- MonadFix Q is not always there
+-------------------------------------------------------------------------------
+
+class MonadFix_ m where
+    mfix_ :: (a -> m a) -> m a
+
+instance MonadFix_ Q where
+    mfix_ k = do
+        m <- runIO newEmptyMVar
+        ans <- runIO (unsafeInterleaveIO (takeMVar m))
+        result <- k ans
+        runIO (putMVar m result)
+        pure result
+    {-# INLINE mfix_ #-}

src/Data/Aeson/Internal/Text.hs

@@ -5,12 +5,11 @@ module Data.Aeson.Internal.Text (
 ) where
 
 import           Data.ByteString                (ByteString)
-import           Data.Text                      (Text, empty)
-
+import qualified Data.Text                      as T
 
 #if MIN_VERSION_text(2,0,0)
 import           Data.Text.Array                (Array (..))
-import           Data.Text.Internal             (Text (..))
+import qualified Data.Text.Internal             as T (Text (..))
 
 import qualified Data.ByteString.Short.Internal as SBS
 
@@ -24,11 +23,11 @@ import qualified Data.Text.Encoding             as TE
 -- | The input is assumed to contain only 7bit ASCII characters (i.e. @< 0x80@).
 --   We use TE.decodeLatin1 here because TE.decodeASCII is currently (text-1.2.4.0)
 --   deprecated and equal to TE.decodeUtf8, which is slower than TE.decodeLatin1.
-unsafeDecodeASCII :: ByteString -> Text
+unsafeDecodeASCII :: ByteString -> T.Text
 
 #if MIN_VERSION_text(2,0,0)
-unsafeDecodeASCII bs = withBS bs $ \_fp len -> if len == 0 then empty else
-  let !(SBS.SBS arr) = SBS.toShort bs in Text (ByteArray arr) 0 len
+unsafeDecodeASCII bs = withBS bs $ \_fp len -> if len == 0 then T.empty else
+  let !(SBS.SBS arr) = SBS.toShort bs in T.Text (ByteArray arr) 0 len
 
 #else
 unsafeDecodeASCII = TE.decodeLatin1