Add primitives to write/read Integers to/from bytearrays

This adds the following new (internal) primitives

{{{#!hs
sizeInBaseInteger :: Integer -> Int# -> Word#

exportInteger :: Integer -> MutableByteArray# s -> Word# -> Int#
                 -> State# s -> (# State# s, Word# #)

importInteger :: ByteArray# -> Word# -> Word# -> Int# -> Integer
}}}

The import/export primitives support selecting most/least significant
byte first order as well as using an offset into the byte-arrays.

See Haddock comments for more details.

Signed-off-by: Herbert Valerio Riedel <hvr@gnu.org>

4 files changed, 142 insertions, 2 deletions

diff --git a/libraries/integer-gmp/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp/GHC/Integer/GMP/Internals.hs
index 127122d081..51727f8e15 100644
--- a/libraries/integer-gmp/GHC/Integer/GMP/Internals.hs
+++ b/libraries/integer-gmp/GHC/Integer/GMP/Internals.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE NoImplicitPrelude #-}
 
-module GHC.Integer.GMP.Internals (Integer(..), gcdInt, gcdInteger, gcdExtInteger, lcmInteger, powInteger, powModInteger, powModSecInteger, recipModInteger, nextPrimeInteger, testPrimeInteger)
+module GHC.Integer.GMP.Internals (Integer(..), gcdInt, gcdInteger, gcdExtInteger, lcmInteger, powInteger, powModInteger, powModSecInteger, recipModInteger, nextPrimeInteger, testPrimeInteger, sizeInBaseInteger, importInteger, exportInteger)
     where
 
 import GHC.Integer.Type
diff --git a/libraries/integer-gmp/GHC/Integer/GMP/Prim.hs b/libraries/integer-gmp/GHC/Integer/GMP/Prim.hs
index beb812cf05..99b5a8ae92 100644
--- a/libraries/integer-gmp/GHC/Integer/GMP/Prim.hs
+++ b/libraries/integer-gmp/GHC/Integer/GMP/Prim.hs
@@ -49,6 +49,10 @@ module GHC.Integer.GMP.Prim (
     nextPrimeInteger#,
     testPrimeInteger#,
 
+    sizeInBaseInteger#,
+    exportInteger#,
+    importInteger#,
+
 #if WORD_SIZE_IN_BITS < 64
     int64ToInteger#,  integerToInt64#,
     word64ToInteger#, integerToWord64#,
@@ -222,6 +226,21 @@ foreign import prim "integer_cmm_testPrimeIntegerzh" testPrimeInteger#
 
 -- |
 --
+foreign import prim "integer_cmm_sizeInBasezh" sizeInBaseInteger#
+  :: Int# -> ByteArray# -> Int# -> Word#
+
+-- |
+--
+foreign import prim "integer_cmm_exportIntegerzh" exportInteger#
+  :: Int# -> ByteArray# -> MutableByteArray# s -> Word# -> Int# -> State# s -> (# State# s, Word# #)
+
+-- |
+--
+foreign import prim "integer_cmm_importIntegerzh" importInteger#
+  :: ByteArray# -> Word# -> Word# -> Int# -> (# Int#, ByteArray# #)
+
+-- |
+--
 foreign import prim "integer_cmm_complementIntegerzh" complementInteger#
   :: Int# -> ByteArray# -> (# Int#, ByteArray# #)
 
diff --git a/libraries/integer-gmp/GHC/Integer/Type.lhs b/libraries/integer-gmp/GHC/Integer/Type.lhs
index a44d29983d..a8d7f09ed3 100644
--- a/libraries/integer-gmp/GHC/Integer/Type.lhs
+++ b/libraries/integer-gmp/GHC/Integer/Type.lhs
@@ -23,7 +23,7 @@ module GHC.Integer.Type where
 
 import GHC.Prim (
     -- Other types we use, convert from, or convert to
-    Int#, Word#, Double#, Float#, ByteArray#,
+    Int#, Word#, Double#, Float#, ByteArray#, MutableByteArray#, State#,
     -- Conversions between those types
     int2Word#, int2Double#, int2Float#, word2Int#,
     -- Operations on Int# that we use for operations on S#
@@ -47,6 +47,7 @@ import GHC.Integer.GMP.Prim (
     testBitInteger#, mul2ExpInteger#, fdivQ2ExpInteger#,
     powInteger#, powModInteger#, powModSecInteger#, recipModInteger#,
     nextPrimeInteger#, testPrimeInteger#,
+    sizeInBaseInteger#, exportInteger#, importInteger#,
 #if WORD_SIZE_IN_BITS < 64
     int64ToInteger#,  integerToInt64#,
     word64ToInteger#, integerToWord64#,
@@ -676,6 +677,76 @@ nextPrimeInteger :: Integer -> Integer
 nextPrimeInteger j@(S# _) = nextPrimeInteger (toBig j)
 nextPrimeInteger (J# s d) = case nextPrimeInteger# s d of (# s', d' #) -> J# s' d'
 
+-- | Compute number of digits (without sign) in given @base@.
+--
+-- It's recommended to avoid calling 'sizeInBaseInteger' for small
+-- integers as this function would currently convert those to big
+-- integers in order to call @mpz_sizeinbase()@.
+--
+-- This function wraps @mpz_sizeinbase()@ which has some
+-- implementation pecularities to take into account:
+--
+-- * @sizeInBaseInteger 0 base = 1@ (see also comment in 'exportInteger').
+--
+-- * This function is only defined if @base >= 2#@ and @base <= 256#@
+--   (Note: the documentation claims that only @base <= 62#@ is
+--   supported, however the actual implementation supports up to base 256).
+--
+-- * If @base@ is a power of 2, the result will be exact. In other
+--   cases (e.g. for @base = 10#@), the result /may/ be 1 digit too large
+--   sometimes.
+--
+-- * @sizeInBaseInteger i 2#@ can be used to determine the most
+--   significant bit of @i@.
+{-# NOINLINE sizeInBaseInteger #-}
+sizeInBaseInteger :: Integer -> Int# -> Word#
+sizeInBaseInteger j@(S# _) b = sizeInBaseInteger (toBig j) b -- TODO
+sizeInBaseInteger (J# s d) b = sizeInBaseInteger# s d b
+
+-- | Dump 'Integer' (without sign) to mutable byte-array in base-256 representation.
+--
+-- The call @exportInteger i mba offset order@ writes
+--
+-- * the 'Integer' @i@
+--
+-- * into the 'MutableByteArray#' @mba@ starting at @offset@
+--
+-- * with most significant byte first if @order@ is @1#@ or least
+--   significant byte first if @order@ is @-1#@, and
+--
+-- * returns number of bytes written.
+--
+-- Use @sizeInBaseInteger i 256#@ to compute the exact number of bytes
+-- written in advance for @i /= 0@. In case of @i == 0@,
+-- 'exportInteger' will write and report zero bytes written, whereas
+-- 'sizeInBaseInteger' report one byte.
+--
+-- It's recommended to avoid calling 'exportInteger' for small
+-- integers as this function would currently convert those to big
+-- integers in order to call @mpz_export()@.
+{-# NOINLINE exportInteger #-}
+exportInteger :: Integer -> MutableByteArray# s -> Word# -> Int# -> State# s -> (# State# s, Word# #)
+exportInteger j@(S# _) mba o e = exportInteger (toBig j) mba o e -- TODO
+exportInteger (J# s d) mba o e = exportInteger# s d mba o e
+
+-- | Read 'Integer' (without sign) from byte-array in base-256 representation.
+--
+-- The call @importInteger ba offset size order@ reads
+--
+-- * @size@ bytes from the 'ByteArray#' @mba@ starting at @offset@
+--
+-- * with most significant byte first if @order@ is @1#@ or least
+--   significant byte first if @order@ is @-1#@, and
+--
+-- * returns a new 'Integer'
+--
+-- It's recommended to avoid calling 'importInteger' for known to be
+-- small integers as this function currently always returns a big
+-- integer even if it would fit into a small integer.
+{-# NOINLINE importInteger #-}
+importInteger :: ByteArray# -> Word# -> Word# -> Int# -> Integer
+importInteger ba o l e = case importInteger# ba o l e of (# s', d' #) -> J# s' d'
+
 \end{code}
 
 %*********************************************************
diff --git a/libraries/integer-gmp/cbits/gmp-wrappers.cmm b/libraries/integer-gmp/cbits/gmp-wrappers.cmm
index 4db684db09..186cfadadc 100644
--- a/libraries/integer-gmp/cbits/gmp-wrappers.cmm
+++ b/libraries/integer-gmp/cbits/gmp-wrappers.cmm
@@ -56,6 +56,9 @@ import "integer-gmp" __gmpz_powm_sec;
 import "integer-gmp" __gmpz_invert;
 import "integer-gmp" __gmpz_nextprime;
 import "integer-gmp" __gmpz_probab_prime_p;
+import "integer-gmp" __gmpz_sizeinbase;
+import "integer-gmp" __gmpz_import;
+import "integer-gmp" __gmpz_export;
 
 import "integer-gmp" integer_cbits_decodeDouble;
 
@@ -66,6 +69,51 @@ import "integer-gmp" integer_cbits_decodeDouble;
    the case for all the platforms that GHC supports, currently.
    -------------------------------------------------------------------------- */
 
+integer_cmm_importIntegerzh (P_ ba, W_ of, W_ sz, W_ e)
+{
+  P_ p;
+  W_ mp_result1;
+
+again:
+  STK_CHK_GEN_N (SIZEOF_MP_INT);
+  MAYBE_GC(again);
+
+  p = ba + SIZEOF_StgArrWords + of;
+
+  mp_result1 = Sp - SIZEOF_MP_INT;
+
+  ccall __gmpz_init(mp_result1 "ptr");
+  ccall __gmpz_import(mp_result1 "ptr", sz, W_TO_INT(e), W_TO_INT(1), W_TO_INT(0), 0, p "ptr");
+
+  return(TO_W_(MP_INT__mp_size(mp_result1)),
+         MP_INT__mp_d(mp_result1) - SIZEOF_StgArrWords);
+}
+
+integer_cmm_exportIntegerzh (W_ s1, P_ d1, P_ mba, W_ of, W_ e)
+{
+  P_ dst;
+  W_ mp_tmp;
+  W_ cnt_result1;
+
+again:
+  STK_CHK_GEN_N (SIZEOF_MP_INT + SIZEOF_W);
+  MAYBE_GC(again);
+
+  mp_tmp     = Sp - SIZEOF_MP_INT;
+  MP_INT__mp_alloc(mp_tmp) = W_TO_INT(BYTE_ARR_WDS(d1));
+  MP_INT__mp_size(mp_tmp)  = (s1);
+  MP_INT__mp_d(mp_tmp)     = BYTE_ARR_CTS(d1);
+
+  cnt_result1 = Sp - (SIZEOF_MP_INT + SIZEOF_W);
+  W_[cnt_result1] = 0;
+
+  dst = mba + SIZEOF_StgArrWords + of;
+
+  ccall __gmpz_export(dst "ptr", cnt_result1 "ptr", W_TO_INT(e), W_TO_INT(1), W_TO_INT(0), 0, mp_tmp "ptr");
+
+  return (W_[cnt_result1]);
+}
+
 integer_cmm_int2Integerzh (W_ val)
 {
    W_ s, p; /* to avoid aliasing */
@@ -471,6 +519,8 @@ GMP_TAKE1_UL1_RET1(integer_cmm_powIntegerzh,        __gmpz_pow_ui)
 GMP_TAKE1_RET1(integer_cmm_nextPrimeIntegerzh,      __gmpz_nextprime)
 GMP_TAKE1_I1_RETI1(integer_cmm_testPrimeIntegerzh,  __gmpz_probab_prime_p)
 
+GMP_TAKE1_I1_RETI1(integer_cmm_sizeInBasezh,        __gmpz_sizeinbase)
+
 integer_cmm_gcdIntzh (W_ int1, W_ int2)
 {
     W_ r;