From efedb3857053996e8c2f718bc40e1d5351c674cc Mon Sep 17 00:00:00 2001
From: Vincent Hanquez <vincent@snarc.org>
Date: Fri, 22 May 2015 21:36:13 +0100
Subject: [PATCH] [number] add a new routine to generate random number

---
 Crypto/Number/Generate.hs | 62 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 2 deletions(-)

diff --git a/Crypto/Number/Generate.hs b/Crypto/Number/Generate.hs
index 912f569..c5b62cf 100644
--- a/Crypto/Number/Generate.hs
+++ b/Crypto/Number/Generate.hs
@@ -6,7 +6,9 @@
 -- Portability : Good
 
 module Crypto.Number.Generate
-    ( generateMax
+    ( GenTopPolicy(..)
+    , generate
+    , generateMax
     , generateBetween
     , generateOfSize
     , generateBits
@@ -16,11 +18,67 @@ import           Crypto.Internal.Imports
 import           Crypto.Number.Basic
 import           Crypto.Number.Serialize
 import           Crypto.Random.Types
-import           Data.Bits ((.|.), (.&.), shiftR)
+import           Control.Monad (when)
+import           Foreign.Ptr
+import           Foreign.Storable
+import           Data.Bits ((.|.), (.&.), shiftL, shiftR, complement)
 import           Crypto.Internal.ByteArray (Bytes, ScrubbedBytes)
 import qualified Crypto.Internal.ByteArray as B
 
 
+-- | Top bits policy when generating a number
+data GenTopPolicy =
+      SetHighest    -- ^ set the highest bit
+    | SetTwoHighest -- ^ set the two highest bit
+    deriving (Show,Eq)
+
+-- | Generate a number for a specific size of bits,
+-- and optionaly set bottom and top bits
+--
+-- If the top bit policy is 'Nothing', then nothing is
+-- done on the highest bit (it's whatever the random generator set).
+--
+-- If @generateOdd is set to 'True', then the number generated
+-- is guaranteed to be odd. Otherwise it will be whatever is generated
+--
+generate :: MonadRandom m
+         => Int                -- ^ number of bits
+         -> Maybe GenTopPolicy -- ^ top bit policy
+         -> Bool               -- ^ force the number to be odd
+         -> m Integer
+generate bits genTopPolicy generateOdd
+    | bits <= 0 = return 0
+    | otherwise = os2ip . tweak <$> getRandomBytes bytes
+  where
+    tweak :: ScrubbedBytes -> ScrubbedBytes
+    tweak orig =
+        case (genTopPolicy, generateOdd) of
+            (Nothing       , False) -> orig
+            (Nothing       , True ) -> B.copyAndFreeze orig $ \p -> (p `plusPtr` (bytes-1)) |= 0x1
+            (Just topPolicy, _    ) -> B.copyAndFreeze orig $ \p0 -> do
+                let p1   = p0 `plusPtr` 1
+                    pEnd = p0 `plusPtr` (bytes - 1)
+                case topPolicy of
+                    SetHighest                -> p0 |= (1 `shiftL` bit)
+                    SetTwoHighest | bit == 0  -> do p0 $= 0x1
+                                                    p1 |= 0x80
+                                  | otherwise -> p0 |= (0x3 `shiftL` (bit - 1))
+                p0 &= (complement $ mask)
+                when generateOdd (pEnd |= 0x1)
+
+    ($=) :: Ptr Word8 -> Word8 -> IO ()
+    ($=) p w = poke p w
+
+    (|=) :: Ptr Word8 -> Word8 -> IO ()
+    (|=) p w = peek p >>= \v -> poke p (v .|. w)
+
+    (&=) :: Ptr Word8 -> Word8 -> IO ()
+    (&=) p w = peek p >>= \v -> poke p (v .&. w)
+
+    bytes = (bits + 7) `div` 8;
+    bit   = (bits - 1) `mod` 8;
+    mask  = 0xff `shiftL` (bit + 1);
+
 -- | generate a positive integer x, s.t. 0 <= x < m
 generateMax :: MonadRandom m => Integer -> m Integer
 generateMax 1 = return 0