scientific-notation-0.1.6.1: bench/Main.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad.ST (runST)
import Control.Monad.ST.Run (runPrimArrayST)
import Data.ByteString.Internal (ByteString (PS))
import Data.Primitive (ByteArray (..), PrimArray, SmallArray)
import Data.Word (Word16)
import GHC.ForeignPtr (ForeignPtr (ForeignPtr), ForeignPtrContents (PlainPtr))
import Gauge (bench, bgroup, whnf)
import Gauge.Main (defaultMain)

import qualified Data.Aeson.Parser as Aeson
import qualified Data.Attoparsec.ByteString.Char8 as Atto
import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Parser as P
import qualified Data.Bytes.Parser.Latin as Latin
import qualified Data.Bytes.Text.Ascii as Ascii
import qualified Data.Primitive as PM
import qualified GHC.Exts as Exts

import qualified Data.Number.Scientific as SCI
import qualified Data.Scientific as SlowSci

main :: IO ()
main =
  defaultMain
    [ bgroup
        "scientific-notation"
        [ bgroup
            "parser"
            [ bench
                "ten-small"
                (whnf (\b -> P.parseByteArray decodeTen b) tenSmall)
            , bench
                "ten-large"
                (whnf (\b -> P.parseByteArray decodeTen b) tenLarge)
            ]
        , bgroup
            "conversion"
            [ bench
                "twenty-word16"
                (whnf (\b -> convertArray16 b) twentyFastSci)
            ]
        ]
    , bgroup
        "scientific"
        [ bgroup
            "parser"
            [ bench "ten-small" $
                whnf
                  ( \b ->
                      Atto.parseOnly
                        (aesonDecodeN 10 [])
                        (fromPinned b)
                  )
                  tenSmall
            , bench "ten-large" $
                whnf
                  ( \b ->
                      Atto.parseOnly
                        (aesonDecodeN 10 [])
                        (fromPinned b)
                  )
                  tenLarge
            ]
        , bgroup
            "conversion"
            [ bench
                "twenty-word16"
                (whnf (\b -> convertSlowArray16 b) twentySlowSci)
            ]
        ]
    ]

-- TODO: In the test suite, we should confirm that parsing this
-- actually succeeds. We intentionally avoid leading plus signs
-- here so that we can compare against aeson.
tenSmall :: ByteArray
tenSmall =
  pin $
    Bytes.toByteArray $
      Ascii.fromString $
        concat
          [ ",4256"
          , ",-125e14"
          , ",5.000006"
          , ",1e100"
          , ",-13.25E-100"
          , ",-653467618"
          , ",-17e+6"
          , ",9999.001"
          , ",0000.002"
          , ",0000.002E1"
          ]

-- TODO: In the test suite, we should confirm that parsing this
-- actually succeeds. We intentionally avoid leading plus signs
-- here so that we can compare against aeson.
tenLarge :: ByteArray
tenLarge =
  pin $
    Bytes.toByteArray $
      Ascii.fromString $
        concat
          [ ",4221465241250205246754620201240240201451991999956"
          , ",242422432499393113113131313131533753.02031243210e13432"
          , ",-0.999999999999999999999999999999999999"
          , ",4.46246246526345643246256423645246224e100"
          , ",42463523462.46246243246256423645246224E24625"
          , ",-82463523462.56246243246256423645246224e-24625"
          , ",82463523462.56246243246256423645246224e+24625"
          , ",-201.562462432462564236452462240240420"
          , ",-0.777777777777777777777777777777777e-777"
          , ",0.987777777777777777777777777777777e-42"
          ]

-- All of these can fit inside a Word16.
twentyPairs :: SmallArray (Int, Int)
twentyPairs =
  Exts.fromList
    [ (2336, 0)
    , (43265, 0)
    , (17, 0)
    , (24, 3)
    , (1, 4)
    , (25, 0)
    , (0, 0)
    , (1900, 0)
    , (65, 0)
    , (1100, 0)
    , (5, 3)
    , (0, 0)
    , (1600, 0)
    , (1500, 0)
    , (2000, 0)
    , (62, 2)
    , (500, 0)
    , (670, 0)
    , (1100, 0)
    , (65500, 0)
    ]

twentyFastSci :: SmallArray SCI.Scientific
twentyFastSci = fmap (uncurry SCI.small) twentyPairs

twentySlowSci :: SmallArray SlowSci.Scientific
twentySlowSci =
  fmap
    (\(x, y) -> SlowSci.scientific (fromIntegral x) y)
    twentyPairs

aesonDecodeN :: Int -> [SlowSci.Scientific] -> Atto.Parser [SlowSci.Scientific]
aesonDecodeN !ix !acc =
  if ix > 0
    then do
      _ <- Atto.char ','
      !num <- Aeson.scientific
      aesonDecodeN (ix - 1) (num : acc)
    else pure acc

decodeTen :: P.Parser () s (SmallArray SCI.Scientific)
decodeTen = do
  arr <- P.effect (PM.newSmallArray 10 errorThunk)
  let go !ix =
        if ix >= 0
          then do
            Latin.char () ','
            !num <- SCI.parserSignedUtf8Bytes ()
            P.effect (PM.writeSmallArray arr ix num)
            go (ix - 1)
          else P.effect (PM.unsafeFreezeSmallArray arr)
  go 9

convertArray16 ::
  SmallArray SCI.Scientific ->
  PrimArray Word16
convertArray16 xs = runPrimArrayST $ do
  let len = PM.sizeofSmallArray xs
  ws <- PM.newPrimArray len
  let go !ix =
        if ix >= 0
          then case SCI.toWord16 (PM.indexSmallArray xs ix) of
            Nothing -> error "convertArray16: bad number"
            Just (r :: Word16) -> do
              PM.writePrimArray ws ix r
              go (ix - 1)
          else PM.unsafeFreezePrimArray ws
  go (len - 1)

convertSlowArray16 ::
  SmallArray SlowSci.Scientific ->
  PrimArray Word16
convertSlowArray16 xs = runPrimArrayST $ do
  let len = PM.sizeofSmallArray xs
  ws <- PM.newPrimArray len
  let go !ix =
        if ix >= 0
          then case SlowSci.toBoundedInteger (PM.indexSmallArray xs ix) of
            Nothing -> error "convertArray16: bad number"
            Just (r :: Word16) -> do
              PM.writePrimArray ws ix r
              go (ix - 1)
          else PM.unsafeFreezePrimArray ws
  go (len - 1)

errorThunk :: a
{-# NOINLINE errorThunk #-}
errorThunk = error "scientific:benchmark error"

-- Convert a pinned immutable byte array to a bytestring.
fromPinned :: ByteArray -> ByteString
{-# INLINE fromPinned #-}
fromPinned (ByteArray arr#) =
  PS
    (ForeignPtr (Exts.byteArrayContents# arr#) (PlainPtr (Exts.unsafeCoerce# arr#)))
    0
    (Exts.I# (Exts.sizeofByteArray# arr#))

pin :: ByteArray -> ByteArray
pin src = runST $ do
  let len = PM.sizeofByteArray src
  dst <- PM.newByteArray len
  PM.copyByteArray dst 0 src 0 len
  PM.unsafeFreezeByteArray dst