improve Natural -> ByteString

Author: janmasrovira

bench2/Benchmark/Nockma/Encoding/Natural.hs

@@ -0,0 +1,40 @@
+module Benchmark.Nockma.Encoding.Natural where
+
+import Juvix.Prelude
+import System.Random
+import Test.Tasty.Bench
+
+randomNatural :: IO Natural
+randomNatural = do
+  sg <- getStdGen
+  let numDigitsMin = 100000
+      numDigitsMax = 1000000
+  return (fst (uniformR (10 ^ numDigitsMin, 10 ^ numDigitsMax) sg))
+
+bm :: Benchmark
+bm =
+  let old = padByteString8 . naturalToByteStringOld
+      new = padByteString8 . naturalToByteString
+   in bgroup
+        "Natural -> ByteString"
+        [ env
+            (randomNatural)
+            (\nat -> bench "Old" (nf old nat)),
+          env
+            (randomNatural)
+            (\nat -> bench "New" (nf new nat))
+          -- env
+          --   (randomNatural)
+          --   ( \nat ->
+          --       bench
+          --         "TEST"
+          --         ( nf
+          --             ( \i ->
+          --                 if old i == new i
+          --                   then True
+          --                   else error "wrong"
+          --             )
+          --             nat
+          --         )
+          --   )
+        ]

bench2/Main.hs

@@ -2,12 +2,14 @@ module Main where
 
 import Benchmark.Effect qualified as Effect
 import Benchmark.Nockma qualified as Nockma
+import Benchmark.Nockma.Encoding.Natural as Natural
 import Juvix.Prelude
 import Test.Tasty.Bench
 
 main :: IO ()
 main =
   defaultMain
-    [ Effect.bm,
-      Nockma.bm
+    [ -- Effect.bm,
+      -- Nockma.bm
+      Natural.bm
     ]

src/Juvix/Compiler/Core/Evaluator.hs

@@ -16,7 +16,7 @@ import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Info.NoDisplayInfo
 import Juvix.Compiler.Core.Pretty
 import Juvix.Compiler.Nockma.Encoding qualified as Encoding
-import Juvix.Compiler.Nockma.Encoding.ByteString (byteStringToIntegerLE, naturalToByteStringLELen)
+import Juvix.Compiler.Nockma.Encoding.ByteString (byteStringToIntegerLE)
 import Juvix.Compiler.Nockma.Encoding.Ed25519 qualified as E
 import Juvix.Compiler.Store.Core.Extra qualified as Store
 import Juvix.Data.Field

src/Juvix/Compiler/Nockma/Encoding/ByteString.hs

@@ -1,17 +1,18 @@
+{-# LANGUAGE MagicHash #-}
+
 module Juvix.Compiler.Nockma.Encoding.ByteString where
 
 import Crypto.Hash.SHA256 qualified as SHA256
 import Data.Bit (Bit)
 import Data.Bit qualified as Bit
-import Data.Bits
 import Data.ByteString qualified as BS
 import Data.ByteString.Base64 qualified as Base64
-import Data.ByteString.Builder qualified as BS
 import Juvix.Compiler.Nockma.Encoding.Base
 import Juvix.Compiler.Nockma.Encoding.Effect.BitReader
 import Juvix.Compiler.Nockma.Encoding.Effect.BitWriter
 import Juvix.Compiler.Nockma.Language
 import Juvix.Prelude.Base
+import Juvix.Prelude.Bytes
 
 -- | Encode an atom to little-endian bytes
 atomToByteString :: (NockNatural a, Member (Error (ErrNockNatural a)) r) => Atom a -> Sem r ByteString
@@ -30,9 +31,6 @@ byteStringToAtom = fmap mkEmptyAtom . fromNatural . byteStringToNatural
 byteStringToNatural :: ByteString -> Natural
 byteStringToNatural = fromInteger . byteStringToIntegerLE
 
-naturalToByteString :: Natural -> ByteString
-naturalToByteString = naturalToByteStringLE
-
 naturalToBase64 :: Natural -> Text
 naturalToBase64 = decodeUtf8 . Base64.encode . naturalToByteString
 
@@ -58,18 +56,6 @@ byteStringToIntegerLEChunked = foldr' go 0 . map (first byteStringChunkToInteger
     byteStringChunkToInteger :: ByteString -> Integer
     byteStringChunkToInteger = BS.foldr' (\b acc -> acc `shiftL` 8 .|. fromIntegral b) 0
 
--- | TODO: this is quadratic (`shiftR` is O(n))
-naturalToByteStringLE :: Natural -> ByteString
-naturalToByteStringLE = BS.toStrict . BS.toLazyByteString . go
-  where
-    go :: Natural -> BS.Builder
-    go = \case
-      0 -> mempty
-      n -> BS.word8 (fromIntegral n) <> go (n `shiftR` 8)
-
-naturalToByteStringLELen :: Int -> Natural -> ByteString
-naturalToByteStringLELen len = padByteString len . naturalToByteStringLE
-
 textToNatural :: Text -> Natural
 textToNatural = byteStringToNatural . encodeUtf8
 
@@ -90,12 +76,6 @@ mkEmptyAtom x =
       _atom = x
     }
 
--- | Pad a ByteString with zeros up to a specified length
-padByteString :: Int -> ByteString -> ByteString
-padByteString n bs
-  | BS.length bs >= n = bs
-  | otherwise = BS.append bs (BS.replicate (n - BS.length bs) 0)
-
 vectorBitsToInteger :: Bit.Vector Bit -> Integer
 vectorBitsToInteger = byteStringToIntegerLEChunked . vectorBitsToByteString
 

src/Juvix/Prelude.hs

@@ -2,6 +2,7 @@ module Juvix.Prelude
   ( module Juvix.Prelude.Base,
     module Juvix.Prelude.Lens,
     module Juvix.Prelude.Generic,
+    module Juvix.Prelude.Bytes,
     module Juvix.Prelude.Trace,
     module Juvix.Prelude.Path,
     module Juvix.Prelude.Prepath,
@@ -11,6 +12,7 @@ where
 
 import Juvix.Data
 import Juvix.Prelude.Base
+import Juvix.Prelude.Bytes
 import Juvix.Prelude.Generic
 import Juvix.Prelude.Lens
 import Juvix.Prelude.Path

src/Juvix/Prelude/Base/Foundation.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE MagicHash #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 {-# OPTIONS_GHC -Wno-redundant-constraints #-}
 
@@ -71,6 +72,7 @@ module Juvix.Prelude.Base.Foundation
     module Text.Read,
     module Text.Show,
     module Text.Show.Unicode,
+    module Data.Bits,
     Data,
     Text,
     pack,
@@ -132,6 +134,7 @@ import Control.Monad.Zip
 import Data.Array qualified as Array
 import Data.Bifunctor hiding (first, second)
 import Data.Bitraversable
+import Data.Bits hiding (And, shift)
 import Data.Bool
 import Data.ByteString (ByteString)
 import Data.Char
@@ -178,6 +181,7 @@ import Data.Map.Strict (Map)
 import Data.Maybe
 import Data.Monoid
 import Data.Ord
+import Data.Primitive.ByteArray qualified as GHCByteArray
 import Data.Semigroup (Semigroup, sconcat, (<>))
 import Data.Serialize (Serialize)
 import Data.Serialize as Serial
@@ -264,6 +268,8 @@ type TextBuilder = LazyText.Builder
 
 type GHCType = GHC.Type
 
+type GHCByteArray = GHCByteArray.ByteArray
+
 type GHCConstraint = GHC.Constraint
 
 type LazyHashMap = LazyHashMap.HashMap
@@ -328,6 +334,9 @@ prime name = case Text.splitOn "'" name of
   [name', num] -> name' <> "'" <> maybe (num <> "'") (show . (+ 1)) (Text.readMaybe (unpack num) :: Maybe Word)
   _ -> name <> "'"
 
+divisibleBy :: (Integral a) => a -> a -> Bool
+divisibleBy a b = a `mod` b == 0
+
 freshName :: HashSet Text -> Text -> Text
 freshName names name | HashSet.member name names = freshName names (prime name)
 freshName _ name = name

src/Juvix/Prelude/Bytes.hs

@@ -0,0 +1,83 @@
+{-# LANGUAGE MagicHash #-}
+
+module Juvix.Prelude.Bytes where
+
+import Data.ByteString qualified as BS
+import Data.ByteString.Builder qualified as BS
+import Data.Primitive.ByteArray qualified as GHCByteArray
+import Data.Primitive.Types
+import GHC.Exts (Word (W#))
+import Juvix.Prelude.Base.Foundation
+
+byteArrayToList :: (Prim a) => GHCByteArray -> [a]
+byteArrayToList = GHCByteArray.foldrByteArray (:) mempty
+
+byteArrayToBytes :: GHCByteArray -> [Word8]
+byteArrayToBytes = byteArrayToList
+
+-- | It assumes bytesPerWord == 8
+naturalToWord64 :: Natural -> [Word64]
+naturalToWord64 = \case
+  NS w -> [fromIntegral (W# w)]
+  NB b -> byteArrayToList (GHCByteArray.ByteArray b)
+
+naturalToBytes :: Natural -> [Word8]
+naturalToBytes = \case
+  NS w -> wordToBytes (W# w)
+  NB b -> byteArrayToBytes (GHCByteArray.ByteArray b)
+
+-- | Pad a ByteString with zeros up to the smallest length such that is
+-- divisible by the given arg
+padByteStringMod :: Int -> ByteString -> ByteString
+padByteStringMod align bs =
+  let (d, m) = divMod (BS.length bs) align
+   in if
+          | m == 0 -> bs
+          | otherwise -> padByteString ((d + 1) * align) bs
+
+padByteString8 :: ByteString -> ByteString
+padByteString8 = padByteStringMod 8
+
+-- | Pad a ByteString with zeros up to a specified length
+padByteString :: Int -> ByteString -> ByteString
+padByteString n bs
+  | BS.length bs >= n = bs
+  | otherwise = BS.append bs (BS.replicate (n - BS.length bs) 0)
+
+naturalToByteStringLELen :: Int -> Natural -> ByteString
+naturalToByteStringLELen len = padByteString len . naturalToByteStringLE
+
+naturalToByteStringOld :: Natural -> ByteString
+naturalToByteStringOld = naturalToByteStringLE
+
+-- | TODO: this is quadratic (`shiftR` is O(n))
+naturalToByteStringLE :: Natural -> ByteString
+naturalToByteStringLE = BS.toStrict . BS.toLazyByteString . go
+  where
+    go :: Natural -> BS.Builder
+    go = \case
+      0 -> mempty
+      n -> BS.word8 (fromIntegral n) <> go (n `shiftR` 8)
+
+-- | Little endian
+naturalToByteStringTest :: Bool -> Natural -> ByteString
+naturalToByteStringTest word64 n
+  -- most common case
+  | word64 && 8 == bytesPerWord =
+      let w :: [Word64] = naturalToWord64 n
+       in build (mconcat (map BS.word64LE w))
+  | otherwise = BS.pack (naturalToBytes n)
+  where
+    build :: BS.Builder -> ByteString
+    build = BS.toStrict . BS.toLazyByteString
+
+naturalToByteString :: Natural -> ByteString
+naturalToByteString = naturalToByteStringTest True
+
+-- | Little endian
+wordToBytes :: Word -> [Word8]
+wordToBytes w = [fromIntegral (w `shiftR` (i * 8)) | i <- [0 .. bytesPerWord - 1]]
+
+-- | Platform dependent
+bytesPerWord :: Int
+bytesPerWord = sizeOf (impossible :: Word)