fixed

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

Limb.hs (10106B)


      1 {-# OPTIONS_HADDOCK prune #-}
      2 {-# LANGUAGE BangPatterns #-}
      3 {-# LANGUAGE MagicHash #-}
      4 {-# LANGUAGE NumericUnderscores #-}
      5 {-# LANGUAGE UnboxedSums #-}
      6 {-# LANGUAGE UnboxedTuples #-}
      7 {-# LANGUAGE ViewPatterns #-}
      8 {-# LANGUAGE UnliftedNewtypes #-}
      9 
     10 -- |
     11 -- Module: Data.Word.Limb
     12 -- Copyright: (c) 2025 Jared Tobin
     13 -- License: MIT
     14 -- Maintainer: Jared Tobin <jared@ppad.tech>
     15 --
     16 -- The primitive 'Limb' type, as well as operations on it.
     17 --
     18 -- All operations run in constant time with respect to inputs, unless
     19 -- specifically indicated otherwise.
     20 
     21 module Data.Word.Limb (
     22   -- * Limb
     23     Limb(..)
     24   , render
     25 
     26   -- * Bit manipulation and representation
     27   , and#
     28   , or#
     29   , not#
     30   , xor#
     31   , bits#
     32   , shl#
     33   , shl1#
     34   , shr#
     35   , shr1#
     36 
     37   -- * Comparison
     38   , eq#
     39   , ne#
     40   , eq_vartime#
     41   , ne_vartime#
     42   , nonzero#
     43   , lt#
     44   , gt#
     45 
     46   -- * Selection
     47   , select#
     48   , cswap#
     49 
     50   -- * Negation
     51 
     52   , neg#
     53 
     54   -- * Arithmetic
     55   , add_o#
     56   , add_c#
     57   , add_w#
     58   , add_s#
     59 
     60   , sub_b#
     61   , sub_w#
     62   , sub_s#
     63 
     64   , mul_c#
     65   , mul_w#
     66   , mul_s#
     67 
     68   , mac#
     69 
     70   -- * Re-exported
     71   , Word(..)
     72   , Word#
     73   ) where
     74 
     75 import qualified Data.Bits as B
     76 import qualified Data.Choice as C
     77 import GHC.Exts (Word#)
     78 import qualified GHC.Exts as Exts
     79 
     80 -- | A 'Limb' is the smallest component of a wider word.
     81 newtype Limb = Limb Word#
     82 
     83 -- | Return a 'Limb' value as a 'String'.
     84 render :: Limb -> String
     85 render (Limb a) = show (Exts.W# a)
     86 
     87 -- comparison -----------------------------------------------------------------
     88 
     89 -- | Constant-time equality comparison.
     90 eq#
     91   :: Limb
     92   -> Limb
     93   -> C.Choice
     94 eq# (Limb a) (Limb b) = C.eq_word# a b
     95 {-# INLINE eq# #-}
     96 
     97 -- | Variable-time equality comparison.
     98 eq_vartime#
     99   :: Limb
    100   -> Limb
    101   -> Bool
    102 eq_vartime# (Limb a) (Limb b) = Exts.isTrue# (Exts.eqWord# a b)
    103 {-# INLINE eq_vartime# #-}
    104 
    105 -- | Constant-time inequality comparison.
    106 ne#
    107   :: Limb
    108   -> Limb
    109   -> C.Choice
    110 ne# a b = C.not (eq# a b)
    111 {-# INLINE ne# #-}
    112 
    113 -- | Variable-time inequality comparison.
    114 ne_vartime#
    115   :: Limb
    116   -> Limb
    117   -> Bool
    118 ne_vartime# a b = not (eq_vartime# a b)
    119 {-# INLINE ne_vartime# #-}
    120 
    121 -- | Constant-time comparison to zero.
    122 nonzero#
    123   :: Limb
    124   -> C.Choice
    125 nonzero# (Limb a) = C.from_word_nonzero# a
    126 {-# INLINE nonzero# #-}
    127 
    128 -- | Constant-time less than comparison.
    129 lt#
    130   :: Limb
    131   -> Limb
    132   -> C.Choice
    133 lt# (Limb a) (Limb b) = C.from_word_lt# a b
    134 {-# INLINE lt# #-}
    135 
    136 -- | Constant-time greater than comparison.
    137 gt#
    138   :: Limb
    139   -> Limb
    140   -> C.Choice
    141 gt# (Limb a) (Limb b) = C.from_word_gt# a b
    142 {-# INLINE gt# #-}
    143 
    144 -- constant-time selection ----------------------------------------------------
    145 
    146 -- | Return a if c is truthy, otherwise return b.
    147 select#
    148   :: Limb     -- ^ a
    149   -> Limb     -- ^ b
    150   -> C.Choice -- ^ c
    151   -> Limb     -- ^ result
    152 select# (Limb a) (Limb b) c = Limb (C.select_word# a b c)
    153 {-# INLINE select# #-}
    154 
    155 -- | Return (# b, a #) if c is truthy, otherwise return (# a, b #).
    156 cswap#
    157   :: Limb             -- ^ a
    158   -> Limb             -- ^ b
    159   -> C.Choice         -- ^ c
    160   -> (# Limb, Limb #) -- ^ result
    161 cswap# (Limb a) (Limb b) c =
    162   let !l = C.select_word# a b c
    163       !r = C.select_word# b a c
    164   in  (# Limb l, Limb r #)
    165 {-# INLINE cswap# #-}
    166 
    167 -- bit manipulation -----------------------------------------------------------
    168 
    169 -- | Bitwise and.
    170 and#
    171   :: Limb -- ^ a
    172   -> Limb -- ^ b
    173   -> Limb -- ^ a & b
    174 and# (Limb a) (Limb b) = Limb (Exts.and# a b)
    175 {-# INLINE and# #-}
    176 
    177 -- | Bitwise or.
    178 or#
    179   :: Limb -- ^ a
    180   -> Limb -- ^ b
    181   -> Limb -- ^ a | b
    182 or# (Limb a) (Limb b) = Limb (Exts.or# a b)
    183 {-# INLINE or# #-}
    184 
    185 -- | Bitwise not.
    186 not#
    187   :: Limb -- ^ a
    188   -> Limb -- ^ not a
    189 not# (Limb a) = Limb (Exts.not# a)
    190 {-# INLINE not# #-}
    191 
    192 -- | Bitwise exclusive or.
    193 xor#
    194   :: Limb -- ^ a
    195   -> Limb -- ^ b
    196   -> Limb -- ^ a ^ b
    197 xor# (Limb a) (Limb b) = Limb (Exts.xor# a b)
    198 {-# INLINE xor# #-}
    199 
    200 -- | Number of bits required to represent this limb.
    201 bits#
    202   :: Limb -- ^ limb
    203   -> Int  -- ^ bits required to represent limb
    204 bits# (Limb a) =
    205   let !_BITS = B.finiteBitSize (0 :: Word)
    206       !zs = B.countLeadingZeros (Exts.W# a)
    207   in  _BITS - zs
    208 {-# INLINE bits# #-}
    209 
    210 -- | Unchecked bit-shift left.
    211 shl#
    212   :: Limb       -- ^ limb
    213   -> Exts.Int#  -- ^ shift amount
    214   -> Limb       -- ^ result
    215 shl# (Limb w) s = Limb (Exts.uncheckedShiftL# w s)
    216 {-# INLINE shl# #-}
    217 
    218 -- | Unchecked bit-shift left by 1, returning the result and carry.
    219 shl1#
    220   :: Limb
    221   -> (# Limb, Limb #)
    222 shl1# (Limb w) =
    223   let !s = case B.finiteBitSize (0 :: Word) of Exts.I# m -> m Exts.-# 1#
    224       !r = Exts.uncheckedShiftL# w 1#
    225       !c = Exts.uncheckedShiftRL# w s
    226   in  (# Limb r, Limb c #)
    227 {-# INLINE shl1# #-}
    228 
    229 -- | Unchecked logical bit-shift right.
    230 shr#
    231   :: Limb       -- ^ limb
    232   -> Exts.Int#  -- ^ shift amount
    233   -> Limb       -- ^ result
    234 shr# (Limb w) s = Limb (Exts.uncheckedShiftRL# w s)
    235 {-# INLINE shr# #-}
    236 
    237 -- | Unchecked logical bit-shift right by 1, returning the result and carry.
    238 shr1#
    239   :: Limb
    240   -> (# Limb, Limb #)
    241 shr1# (Limb w) =
    242   let !s = case B.finiteBitSize (0 :: Word) of Exts.I# m -> m Exts.-# 1#
    243       !r = Exts.uncheckedShiftRL# w 1#
    244       !c = Exts.uncheckedShiftL# w s
    245   in  (# Limb r, Limb c #)
    246 {-# INLINE shr1# #-}
    247 
    248 -- negation -------------------------------------------------------------------
    249 
    250 -- | Wrapping (two's complement) negation.
    251 neg#
    252   :: Limb
    253   -> Limb
    254 neg# (Limb x) = Limb (Exts.plusWord# (Exts.not# x) 1##)
    255 {-# INLINE neg# #-}
    256 
    257 -- addition -------------------------------------------------------------------
    258 
    259 -- | Overflowing addition, computing augend + addend, returning the
    260 --   sum and carry.
    261 add_o#
    262   :: Limb             -- ^ augend
    263   -> Limb             -- ^ addend
    264   -> (# Limb, Limb #) -- ^ (# sum, carry #)
    265 add_o# (Limb a) (Limb b) = case Exts.plusWord2# a b of
    266   (# c, s #) -> (# Limb s, Limb c #)
    267 {-# INLINE add_o# #-}
    268 
    269 -- | Carrying addition, computing augend + addend + carry, returning
    270 --   the sum and new carry.
    271 add_c#
    272   :: Limb             -- ^ augend
    273   -> Limb             -- ^ addend
    274   -> Limb             -- ^ carry
    275   -> (# Limb, Limb #) -- ^ (# sum, new carry #)
    276 add_c# (Limb a) (Limb b) (Limb c) =
    277   let !(# c0, s0 #) = Exts.plusWord2# a b
    278       !(# c1,  s #) = Exts.plusWord2# s0 c
    279   in  (# Limb s, Limb (Exts.or# c0 c1) #)
    280 {-# INLINE add_c# #-}
    281 
    282 -- | Wrapping addition, computing augend + addend, returning the sum
    283 --   (discarding overflow).
    284 add_w#
    285   :: Limb -- ^ augend
    286   -> Limb -- ^ addend
    287   -> Limb -- ^ sum
    288 add_w# (Limb a) (Limb b) = Limb (Exts.plusWord# a b)
    289 {-# INLINE add_w# #-}
    290 
    291 -- | Saturating addition, computing augend + addend, returning the
    292 --   sum (clamping to the maximum representable value in the case of
    293 --   overflow).
    294 add_s#
    295   :: Limb
    296   -> Limb
    297   -> Limb
    298 add_s# (Limb a) (Limb b) =
    299   let !(# c, s #) = Exts.plusWord2# a b
    300   in  Limb (C.select_word# s (Exts.not# 0##) (C.from_word_nonzero# c))
    301 {-# INLINE add_s# #-}
    302 
    303 -- subtraction ----------------------------------------------------------------
    304 
    305 -- | Borrowing subtraction, computing minuend - (subtrahend + borrow),
    306 --   returning the difference and new borrow mask.
    307 sub_b#
    308   :: Limb              -- ^ minuend
    309   -> Limb              -- ^ subtrahend
    310   -> Limb              -- ^ borrow
    311   -> (# Limb, Limb #)  -- ^ (# difference, new borrow #)
    312 sub_b# (Limb m) (Limb n) (Limb a) =
    313   let !s = case B.finiteBitSize (0 :: Word) of Exts.I# bs -> bs Exts.-# 1#
    314       !b = Exts.uncheckedShiftRL# a s
    315       !(# d0, b0 #) = Exts.subWordC# m n
    316       !(#  d, b1 #) = Exts.subWordC# d0 b
    317       !c = Exts.int2Word# (Exts.negateInt# (Exts.orI# b0 b1))
    318   in  (# Limb d, Limb c #)
    319 {-# INLINE sub_b# #-}
    320 
    321 -- | Saturating subtraction, computing minuend - subtrahend, returning the
    322 --   difference (and clamping to zero in the case of underflow).
    323 sub_s#
    324   :: Limb -- ^ minuend
    325   -> Limb -- ^ subtrahend
    326   -> Limb -- ^ difference
    327 sub_s# (Limb m) (Limb n) =
    328   let !(# d, b #) = Exts.subWordC# m n
    329       !borrow = C.from_bit# (Exts.int2Word# b)
    330   in  Limb (C.select_word# d 0## borrow)
    331 {-# INLINE sub_s# #-}
    332 
    333 -- | Wrapping subtraction, computing minuend - subtrahend, returning the
    334 --   difference (and discarding underflow).
    335 sub_w#
    336   :: Limb -- ^ minuend
    337   -> Limb -- ^ subtrahend
    338   -> Limb -- ^ difference
    339 sub_w# (Limb m) (Limb n) = Limb (Exts.minusWord# m n)
    340 {-# INLINE sub_w# #-}
    341 
    342 -- multiplication -------------------------------------------------------------
    343 
    344 -- | Widening multiplication, returning low and high words of the product.
    345 mul_c#
    346   :: Limb             -- ^ multiplicand
    347   -> Limb             -- ^ multiplier
    348   -> (# Limb, Limb #) -- ^ (# low, high #) product
    349 mul_c# (Limb a) (Limb b) =
    350   let !(# h, l #) = Exts.timesWord2# a b
    351   in  (# Limb l, Limb h #)
    352 {-# INLINE mul_c# #-}
    353 
    354 -- | Wrapping multiplication, returning only the low word of the product.
    355 mul_w#
    356   :: Limb -- ^ multiplicand
    357   -> Limb -- ^ multiplier
    358   -> Limb -- ^ low word of product
    359 mul_w# (Limb a) (Limb b) = Limb (Exts.timesWord# a b)
    360 {-# INLINE mul_w# #-}
    361 
    362 -- | Saturating multiplication, returning only the low word of the product,
    363 --   and clamping to the maximum value in the case of overflow.
    364 mul_s#
    365   :: Limb -- ^ multiplicand
    366   -> Limb -- ^ multiplier
    367   -> Limb -- ^ clamped low word of product
    368 mul_s# (Limb a) (Limb b) =
    369   let !(# h, l #) = Exts.timesWord2# a b
    370   in  Limb (C.select_word# l (Exts.not# 0##) (C.from_word_nonzero# h))
    371 {-# INLINE mul_s# #-}
    372 
    373 -- | Multiply-add-carry, computing a * b + m + c, returning the
    374 --   result along with the new carry.
    375 mac#
    376   :: Limb              -- ^ a (multiplicand)
    377   -> Limb              -- ^ b (multiplier)
    378   -> Limb              -- ^ m (addend)
    379   -> Limb              -- ^ c (carry)
    380   -> (# Limb, Limb #)  -- ^ a * b + m + c
    381 mac# (Limb a) (Limb b) (Limb m) (Limb c) =
    382     let !(# h, l #) = Exts.timesWord2# a b
    383         !(# l_0, h_0 #) = wadd_w# (# l, h #) m
    384         !(# d, l_1 #) = Exts.plusWord2# l_0 c
    385         !h_1 = Exts.plusWord# h_0 d
    386     in  (# Limb l_1, Limb h_1 #)
    387   where
    388     -- wide wrapping addition
    389     wadd_w# :: (# Word#, Word# #) -> Word# -> (# Word#, Word# #)
    390     wadd_w# (# x_lo, x_hi #) y_lo =
    391       let !(# c0, s0 #) = Exts.plusWord2# x_lo y_lo
    392           !(# _, s1 #) = Exts.plusWord2# x_hi c0
    393       in  (# s0, s1 #)
    394     {-# INLINE wadd_w# #-}
    395 {-# INLINE mac# #-}
    396