fixed

Pure Haskell large fixed-width integers and Montgomery arithmetic.
git clone git://git.ppad.tech/fixed.git
Log | Files | Refs | README | LICENSE

Wide.hs (7064B)


      1 {-# LANGUAGE BangPatterns #-}
      2 {-# LANGUAGE MagicHash #-}
      3 {-# LANGUAGE NumericUnderscores #-}
      4 {-# LANGUAGE PatternSynonyms #-}
      5 {-# LANGUAGE ViewPatterns #-}
      6 {-# LANGUAGE UnboxedSums #-}
      7 {-# LANGUAGE UnboxedTuples #-}
      8 {-# LANGUAGE UnliftedNewtypes #-}
      9 
     10 -- |
     11 -- Module: Data.Word.Wide
     12 -- Copyright: (c) 2025 Jared Tobin
     13 -- License: MIT
     14 -- Maintainer: Jared Tobin <jared@ppad.tech>
     15 --
     16 -- Wide words, consisting of two 'Limb's.
     17 
     18 module Data.Word.Wide (
     19   -- * Wide Words
     20     Wide(..)
     21 
     22   -- * Construction, Conversion
     23   , wide
     24   , to_vartime
     25   , from_vartime
     26 
     27   -- * Bit Manipulation
     28   , or
     29   , or#
     30   , and
     31   , and#
     32   , xor
     33   , xor#
     34   , not
     35   , not#
     36 
     37   -- * Comparison
     38   , eq_vartime
     39 
     40   -- * Arithmetic
     41   , add
     42   , add_o
     43   , sub
     44   , mul
     45   , neg
     46 
     47   -- * Unboxed Arithmetic
     48   , add_o#
     49   , add_w#
     50   , sub_b#
     51   , sub_w#
     52   , mul_w#
     53   , neg#
     54   ) where
     55 
     56 import Control.DeepSeq
     57 import Data.Bits ((.|.), (.&.), (.<<.), (.>>.))
     58 import qualified Data.Bits as B
     59 import qualified Data.Choice as C
     60 import Data.Word.Limb (Limb(..))
     61 import qualified Data.Word.Limb as L
     62 import GHC.Exts
     63 import Prelude hiding (div, mod, or, and, not, quot, rem, recip)
     64 
     65 -- utilities ------------------------------------------------------------------
     66 
     67 fi :: (Integral a, Num b) => a -> b
     68 fi = fromIntegral
     69 {-# INLINE fi #-}
     70 
     71 -- wide words -----------------------------------------------------------------
     72 
     73 pattern Limb2
     74   :: Word# -> Word#
     75   -> (# Limb, Limb #)
     76 pattern Limb2 w0 w1 = (# Limb w0, Limb w1 #)
     77 {-# COMPLETE Limb2 #-}
     78 
     79 -- | Little-endian wide words.
     80 data Wide = Wide !(# Limb, Limb #)
     81 
     82 instance Show Wide where
     83   show = show . from_vartime
     84 
     85 -- | Note that 'fromInteger' necessarily runs in variable time due
     86 --   to conversion from the variable-size, potentially heap-allocated
     87 --   'Integer' type.
     88 instance Num Wide where
     89   (+) = add
     90   (-) = sub
     91   (*) = mul
     92   abs = id
     93   fromInteger = to_vartime
     94   negate = neg
     95   signum (Wide (# l0, l1 #)) =
     96     let !(Limb l) = l0 `L.or#` l1
     97         !n = C.from_word_nonzero# l
     98         !b = C.to_word# n
     99     in  Wide (Limb2 b 0##)
    100 
    101 instance NFData Wide where
    102   rnf (Wide a) = case a of (# _, _ #) -> ()
    103 
    104 -- construction / conversion --------------------------------------------------
    105 
    106 -- | Construct a 'Wide' word from low and high 'Word's.
    107 wide :: Word -> Word -> Wide
    108 wide (W# l) (W# h) = Wide (# Limb l, Limb h #)
    109 
    110 -- | Convert an 'Integer' to a 'Wide' word.
    111 --
    112 --   >>> to_vartime 1
    113 --   1
    114 to_vartime :: Integer -> Wide
    115 to_vartime n =
    116   let !size = B.finiteBitSize (0 :: Word)
    117       !mask = fi (maxBound :: Word) :: Integer
    118       !(W# w0) = fi (n .&. mask)
    119       !(W# w1) = fi ((n .>>. size) .&. mask)
    120   in  Wide (# Limb w0, Limb w1 #)
    121 
    122 -- | Convert a 'Wide' word to an 'Integer'.
    123 --
    124 --   >>> from_vartime 1
    125 --   1
    126 from_vartime :: Wide -> Integer
    127 from_vartime (Wide (# Limb a, Limb b #)) =
    128       fi (W# b) .<<. (B.finiteBitSize (0 :: Word))
    129   .|. fi (W# a)
    130 
    131 -- comparison -----------------------------------------------------------------
    132 
    133 -- | Compare 'Wide' words for equality in variable time.
    134 eq_vartime :: Wide -> Wide -> Bool
    135 eq_vartime (Wide (# Limb a0, Limb b0 #)) (Wide (# Limb a1, Limb b1 #)) =
    136   isTrue# (andI# (eqWord# a0 a1) (eqWord# b0 b1))
    137 
    138 -- bits -----------------------------------------------------------------------
    139 
    140 or_w# :: (# Limb, Limb #) -> (# Limb, Limb #) -> (# Limb, Limb #)
    141 or_w# (# a0, a1 #) (# b0, b1 #) = (# L.or# a0 b0, L.or# a1 b1 #)
    142 {-# INLINE or_w# #-}
    143 
    144 or :: Wide -> Wide -> Wide
    145 or (Wide a) (Wide b) = Wide (or_w# a b)
    146 
    147 and_w# :: (# Limb, Limb #) -> (# Limb, Limb #) -> (# Limb, Limb #)
    148 and_w# (# a0, a1 #) (# b0, b1 #) = (# L.and# a0 b0, L.and# a1 b1 #)
    149 {-# INLINE and_w# #-}
    150 
    151 and :: Wide -> Wide -> Wide
    152 and (Wide a) (Wide b) = Wide (and_w# a b)
    153 
    154 xor_w# :: (# Limb, Limb #) -> (# Limb, Limb #) -> (# Limb, Limb #)
    155 xor_w# (# a0, a1 #) (# b0, b1 #) = (# L.xor# a0 b0, L.xor# a1 b1 #)
    156 {-# INLINE xor_w# #-}
    157 
    158 xor :: Wide -> Wide -> Wide
    159 xor (Wide a) (Wide b) = Wide (xor_w# a b)
    160 
    161 not_w# :: (# Limb, Limb #) -> (# Limb, Limb #)
    162 not_w# (# a0, a1 #) = (# L.not# a0, L.not# a1 #)
    163 {-# INLINE not_w# #-}
    164 
    165 not :: Wide -> Wide
    166 not (Wide w) = Wide (not_w# w)
    167 {-# INLINE not #-}
    168 
    169 -- negation -------------------------------------------------------------------
    170 
    171 neg#
    172   :: (# Limb, Limb #) -- ^ argument
    173   -> (# Limb, Limb #) -- ^ (wrapping) additive inverse
    174 neg# w = add_w# (not_w# w) (# Limb 1##, Limb 0## #)
    175 {-# INLINE neg# #-}
    176 
    177 neg
    178   :: Wide -- ^ argument
    179   -> Wide -- ^ (wrapping) additive inverse
    180 neg (Wide w) = Wide (neg# w)
    181 
    182 -- addition, subtraction ------------------------------------------------------
    183 
    184 -- | Overflowing addition, computing 'a + b', returning the sum and a
    185 --   carry bit.
    186 add_o#
    187   :: (# Limb, Limb #)              -- ^ augend
    188   -> (# Limb, Limb #)              -- ^ addend
    189   -> (# (# Limb, Limb #), Limb #)  -- ^ (# sum, carry bit #)
    190 add_o# (# a0, a1 #) (# b0, b1 #) =
    191   let !(# s0, c0 #) = L.add_o# a0 b0
    192       !(# s1, c1 #) = L.add_c# a1 b1 c0
    193   in  (# (# s0, s1 #), c1 #)
    194 {-# INLINE add_o# #-}
    195 
    196 -- | Overflowing addition on 'Wide' words, computing 'a + b', returning
    197 --   the sum and carry.
    198 add_o
    199   :: Wide         -- ^ augend
    200   -> Wide         -- ^ addend
    201   -> (Wide, Word) -- ^ (sum, carry)
    202 add_o (Wide a) (Wide b) =
    203   let !(# s, Limb c #) = add_o# a b
    204   in  (Wide s, W# c)
    205 
    206 -- | Wrapping addition, computing 'a + b'.
    207 add_w#
    208   :: (# Limb, Limb #) -- ^ augend
    209   -> (# Limb, Limb #) -- ^ addend
    210   -> (# Limb, Limb #) -- ^ sum
    211 add_w# a b =
    212   let !(# c, _ #) = add_o# a b
    213   in  c
    214 {-# INLINE add_w# #-}
    215 
    216 -- | Wrapping addition on 'Wide' words, computing 'a + b'.
    217 add :: Wide -> Wide -> Wide
    218 add (Wide a) (Wide b) = Wide (add_w# a b)
    219 
    220 -- | Borrowing subtraction, computing 'a - b' and returning the
    221 --   difference with a borrow mask.
    222 sub_b#
    223   :: (# Limb, Limb #)              -- ^ minuend
    224   -> (# Limb, Limb #)              -- ^ subtrahend
    225   -> (# (# Limb, Limb #), Limb #) -- ^ (# difference, borrow mask #)
    226 sub_b# (# a0, a1 #) (# b0, b1 #) =
    227   let !(# s0, c0 #) = L.sub_b# a0 b0 (Limb 0##)
    228       !(# s1, c1 #) = L.sub_b# a1 b1 c0
    229   in  (# (# s0, s1 #), c1 #)
    230 {-# INLINE sub_b# #-}
    231 
    232 -- | Wrapping subtraction, computing 'a - b'.
    233 sub_w#
    234   :: (# Limb, Limb #) -- ^ minuend
    235   -> (# Limb, Limb #) -- ^ subtrahend
    236   -> (# Limb, Limb #) -- ^ difference
    237 sub_w# a b =
    238   let !(# c, _ #) = sub_b# a b
    239   in  c
    240 {-# INLINE sub_w# #-}
    241 
    242 -- | Wrapping subtraction on 'Wide' words, computing 'a - b'.
    243 sub :: Wide -> Wide -> Wide
    244 sub (Wide a) (Wide b) = Wide (sub_w# a b)
    245 
    246 -- multiplication -------------------------------------------------------------
    247 
    248 -- | Wrapping multiplication, computing 'a b'.
    249 mul_w#
    250   :: (# Limb, Limb #) -- ^ multiplicand
    251   -> (# Limb, Limb #) -- ^ multiplier
    252   -> (# Limb, Limb #) -- ^ product
    253 mul_w# (# a0, a1 #) (# b0, b1 #) =
    254   let !(# p0_lo, p0_hi #) = L.mul_c# a0 b0
    255       !(# p1_lo, _ #) = L.mul_c# a0 b1
    256       !(# p2_lo, _ #) = L.mul_c# a1 b0
    257       !(# s0, _ #) = L.add_o# p0_hi p1_lo
    258       !(# s1, _ #) = L.add_o# s0 p2_lo
    259   in  (# p0_lo, s1 #)
    260 {-# INLINE mul_w# #-}
    261 
    262 -- | Wrapping multiplication on 'Wide' words.
    263 mul :: Wide -> Wide -> Wide
    264 mul (Wide a) (Wide b) = Wide (mul_w# a b)
    265