From 6028e958057cf83fadf76f9c671a4e727332f19a Mon Sep 17 00:00:00 2001
From: Vincent Hanquez <vincent@snarc.org>
Date: Mon, 1 Jun 2015 11:50:24 +0100
Subject: [PATCH] [number] cleanup

---
 Crypto/Number/ModArithmetic.hs      |  2 +-
 Crypto/Number/Serialize.hs          | 28 ++++++++++------------------
 Crypto/Number/Serialize/Internal.hs | 27 ++++++++++++---------------
 3 files changed, 23 insertions(+), 34 deletions(-)

diff --git a/Crypto/Number/ModArithmetic.hs b/Crypto/Number/ModArithmetic.hs
index 3d262f3..ed18805 100644
--- a/Crypto/Number/ModArithmetic.hs
+++ b/Crypto/Number/ModArithmetic.hs
@@ -47,7 +47,7 @@ expSafe :: Integer -- ^ base
         -> Integer -- ^ result
 expSafe b e m
     | odd m     = gmpPowModSecInteger b e m `onGmpUnsupported`
-                  (gmpPowModInteger b e m    `onGmpUnsupported`
+                  (gmpPowModInteger b e m   `onGmpUnsupported`
                   exponentiation b e m)
     | otherwise = gmpPowModInteger b e m    `onGmpUnsupported`
                   exponentiation b e m
diff --git a/Crypto/Number/Serialize.hs b/Crypto/Number/Serialize.hs
index b1b096f..9533706 100644
--- a/Crypto/Number/Serialize.hs
+++ b/Crypto/Number/Serialize.hs
@@ -12,29 +12,28 @@ module Crypto.Number.Serialize
     , os2ip
     , i2ospOf
     , i2ospOf_
-    , lengthBytes
     ) where
 
-import           Data.Bits
 import           Data.Word
 import           Foreign.Storable
-import           Crypto.Number.Compat
+import           Crypto.Number.Basic
 import           Crypto.Internal.Compat (unsafeDoIO)
 import qualified Crypto.Number.Serialize.Internal as Internal
 import qualified Crypto.Internal.ByteArray as B
 
 -- | os2ip converts a byte string into a positive integer
 os2ip :: B.ByteArrayAccess ba => ba -> Integer
-os2ip bs = unsafeDoIO $ B.withByteArray bs $ \src -> Internal.os2ip src (B.length bs)
+os2ip bs = unsafeDoIO $! B.withByteArray bs $ \src -> Internal.os2ip src (B.length bs)
 
 -- | i2osp converts a positive integer into a byte string
 --
 -- first byte is MSB (most significant byte), last byte is the LSB (least significant byte)
 i2osp :: B.ByteArray ba => Integer -> ba
-i2osp 0 = B.allocAndFreeze 1 $ \p -> pokeByteOff p 0 (0 :: Word8)
-i2osp m = B.allocAndFreeze sz (\p -> Internal.i2osp m p sz >> return ())
+i2osp 0 = B.unsafeCreate 1 $ \p -> pokeByteOff p 0 (0 :: Word8)
+i2osp m = B.unsafeCreate sz (\p -> Internal.i2osp m p sz >> return ())
   where
-    !sz = lengthBytes m
+    !sz = numBytes m
+{-# NOINLINE i2osp #-}
 
 -- | just like i2osp, but take an extra parameter for size.
 -- if the number is too big to fit in @len bytes, nothing is returned
@@ -43,9 +42,11 @@ i2ospOf :: B.ByteArray ba => Int -> Integer -> Maybe ba
 i2ospOf len m
     | len <= 0  = Nothing
     | m < 0     = Nothing
-    | otherwise = Just $ B.unsafeCreate len $ \p -> Internal.i2ospOf m p len >> return ()
+    | otherwise = Just $ B.unsafeCreate len $ \p -> do
+        b <- Internal.i2ospOf m p len
+        if (b /= len) then error "invalid" else return ()
+{-# NOINLINE i2ospOf #-}
 
---
 -- | just like i2ospOf except that it doesn't expect a failure: i.e.
 -- an integer larger than the number of output bytes requested
 --
@@ -54,12 +55,3 @@ i2ospOf len m
 i2ospOf_ :: B.ByteArray ba => Int -> Integer -> ba
 i2ospOf_ len = maybe (error "i2ospOf_: integer is larger than expected") id . i2ospOf len
 
--- | returns the number of bytes to store an integer with i2osp
---
--- with integer-simple, this function is really slow.
-lengthBytes :: Integer -> Int
-lengthBytes n = gmpSizeInBytes n `onGmpUnsupported` nbBytes n
-  where
-    nbBytes !v
-        | v < 256   = 1
-        | otherwise = 1 + nbBytes (v `shiftR` 8)
diff --git a/Crypto/Number/Serialize/Internal.hs b/Crypto/Number/Serialize/Internal.hs
index b729891..a819e6d 100644
--- a/Crypto/Number/Serialize/Internal.hs
+++ b/Crypto/Number/Serialize/Internal.hs
@@ -33,17 +33,9 @@ i2osp m ptr ptrSz
     | m < 0      = return 0
     | m == 0     = pokeByteOff ptr 0 (0 :: Word8) >> return 1
     | ptrSz < sz = return 0
-    | otherwise  = fillPtr >> return sz
+    | otherwise  = fillPtr ptr sz m >> return sz
   where
-    !sz = numBytes m
-
-    fillPtr = gmpExportInteger m ptr `onGmpUnsupported` export ptr (sz-1) m
-    export p ofs i
-        | ofs == 0  = pokeByteOff p ofs (fromIntegral i :: Word8)
-        | otherwise = do
-            let (i', b) = i `divMod` 256
-            pokeByteOff p ofs (fromIntegral b :: Word8)
-            export p (ofs-1) i'
+    !sz    = numBytes m
 
 -- | Similar to 'i2osp', except it will pad any remaining space with zero.
 i2ospOf :: Integer -> Ptr Word8 -> Int -> IO Int
@@ -51,25 +43,30 @@ i2ospOf m ptr ptrSz
     | ptrSz <= 0 = return 0
     | m < 0      = return 0
     | ptrSz < sz = return 0
-    | otherwise  = (if padSz > 0 then memSet ptr 0 padSz else return ()) >> fillPtr (ptr `plusPtr` padSz) >> return ptrSz
+    | otherwise  = do
+        if padSz > 0 then memSet ptr 0 padSz else return ()
+        fillPtr (ptr `plusPtr` padSz) sz m
+        return ptrSz
   where
     !sz    = numBytes m
     !padSz = ptrSz - sz
 
-    fillPtr p = gmpExportInteger m p `onGmpUnsupported` export p (sz-1) m
-    export p ofs i
+fillPtr :: Ptr Word8 -> Int -> Integer -> IO ()
+fillPtr p sz m = gmpExportInteger m p `onGmpUnsupported` export (sz-1) m
+  where
+    export ofs i
         | ofs == 0  = pokeByteOff p ofs (fromIntegral i :: Word8)
         | otherwise = do
             let (i', b) = i `divMod` 256
             pokeByteOff p ofs (fromIntegral b :: Word8)
-            export p (ofs-1) i'
+            export (ofs-1) i'
 
 -- | transform a big endian binary integer representation pointed by a pointer and a size
 -- into an integer
 os2ip :: Ptr Word8 -> Int -> IO Integer
 os2ip ptr ptrSz
     | ptrSz <= 0 = return 0
-    | otherwise  = {-gmpImportInteger ptrSz ptr `onGmpUnsupported` -} loop 0 0 ptr 
+    | otherwise  = gmpImportInteger ptrSz ptr `onGmpUnsupported` loop 0 0 ptr
   where
     loop :: Integer -> Int -> Ptr Word8 -> IO Integer
     loop !acc i p