fixed

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

Limb.hs (9627B)


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