Limb.hs (10416B)
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 b if c is truthy, otherwise return a. 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 -- 272 -- The carry-in must be 0 or 1; this is not checked. The carry-out 273 -- is always 0 or 1. 274 add_c# 275 :: Limb -- ^ augend 276 -> Limb -- ^ addend 277 -> Limb -- ^ carry 278 -> (# Limb, Limb #) -- ^ (# sum, new carry #) 279 add_c# (Limb a) (Limb b) (Limb c) = 280 let !(# c0, s0 #) = Exts.plusWord2# a b 281 !(# c1, s #) = Exts.plusWord2# s0 c 282 in (# Limb s, Limb (Exts.or# c0 c1) #) 283 {-# INLINE add_c# #-} 284 285 -- | Wrapping addition, computing augend + addend, returning the sum 286 -- (discarding overflow). 287 add_w# 288 :: Limb -- ^ augend 289 -> Limb -- ^ addend 290 -> Limb -- ^ sum 291 add_w# (Limb a) (Limb b) = Limb (Exts.plusWord# a b) 292 {-# INLINE add_w# #-} 293 294 -- | Saturating addition, computing augend + addend, returning the 295 -- sum (clamping to the maximum representable value in the case of 296 -- overflow). 297 add_s# 298 :: Limb 299 -> Limb 300 -> Limb 301 add_s# (Limb a) (Limb b) = 302 let !(# c, s #) = Exts.plusWord2# a b 303 in Limb (C.select_word# s (Exts.not# 0##) (C.from_word_nonzero# c)) 304 {-# INLINE add_s# #-} 305 306 -- subtraction ---------------------------------------------------------------- 307 308 -- | Borrowing subtraction, computing minuend - (subtrahend + borrow), 309 -- returning the difference and new borrow mask. 310 -- 311 -- The borrow-in is interpreted as a mask, of which only the most 312 -- significant bit is examined; a full-word mask (as produced by 313 -- this function) borrows 1. The borrow-out is 0 or a full-word 314 -- mask. 315 sub_b# 316 :: Limb -- ^ minuend 317 -> Limb -- ^ subtrahend 318 -> Limb -- ^ borrow 319 -> (# Limb, Limb #) -- ^ (# difference, new borrow #) 320 sub_b# (Limb m) (Limb n) (Limb a) = 321 let !s = case B.finiteBitSize (0 :: Word) of Exts.I# bs -> bs Exts.-# 1# 322 !b = Exts.uncheckedShiftRL# a s 323 !(# d0, b0 #) = Exts.subWordC# m n 324 !(# d, b1 #) = Exts.subWordC# d0 b 325 !c = Exts.int2Word# (Exts.negateInt# (Exts.orI# b0 b1)) 326 in (# Limb d, Limb c #) 327 {-# INLINE sub_b# #-} 328 329 -- | Saturating subtraction, computing minuend - subtrahend, returning the 330 -- difference (and clamping to zero in the case of underflow). 331 sub_s# 332 :: Limb -- ^ minuend 333 -> Limb -- ^ subtrahend 334 -> Limb -- ^ difference 335 sub_s# (Limb m) (Limb n) = 336 let !(# d, b #) = Exts.subWordC# m n 337 !borrow = C.from_bit# (Exts.int2Word# b) 338 in Limb (C.select_word# d 0## borrow) 339 {-# INLINE sub_s# #-} 340 341 -- | Wrapping subtraction, computing minuend - subtrahend, returning the 342 -- difference (and discarding underflow). 343 sub_w# 344 :: Limb -- ^ minuend 345 -> Limb -- ^ subtrahend 346 -> Limb -- ^ difference 347 sub_w# (Limb m) (Limb n) = Limb (Exts.minusWord# m n) 348 {-# INLINE sub_w# #-} 349 350 -- multiplication ------------------------------------------------------------- 351 352 -- | Widening multiplication, returning low and high words of the product. 353 mul_c# 354 :: Limb -- ^ multiplicand 355 -> Limb -- ^ multiplier 356 -> (# Limb, Limb #) -- ^ (# low, high #) product 357 mul_c# (Limb a) (Limb b) = 358 let !(# h, l #) = Exts.timesWord2# a b 359 in (# Limb l, Limb h #) 360 {-# INLINE mul_c# #-} 361 362 -- | Wrapping multiplication, returning only the low word of the product. 363 mul_w# 364 :: Limb -- ^ multiplicand 365 -> Limb -- ^ multiplier 366 -> Limb -- ^ low word of product 367 mul_w# (Limb a) (Limb b) = Limb (Exts.timesWord# a b) 368 {-# INLINE mul_w# #-} 369 370 -- | Saturating multiplication, returning only the low word of the product, 371 -- and clamping to the maximum value in the case of overflow. 372 mul_s# 373 :: Limb -- ^ multiplicand 374 -> Limb -- ^ multiplier 375 -> Limb -- ^ clamped low word of product 376 mul_s# (Limb a) (Limb b) = 377 let !(# h, l #) = Exts.timesWord2# a b 378 in Limb (C.select_word# l (Exts.not# 0##) (C.from_word_nonzero# h)) 379 {-# INLINE mul_s# #-} 380 381 -- | Multiply-add-carry, computing a * b + m + c, returning the 382 -- result along with the new carry. 383 mac# 384 :: Limb -- ^ a (multiplicand) 385 -> Limb -- ^ b (multiplier) 386 -> Limb -- ^ m (addend) 387 -> Limb -- ^ c (carry) 388 -> (# Limb, Limb #) -- ^ a * b + m + c 389 mac# (Limb a) (Limb b) (Limb m) (Limb c) = 390 let !(# h, l #) = Exts.timesWord2# a b 391 !(# l_0, h_0 #) = wadd_w# (# l, h #) m 392 !(# d, l_1 #) = Exts.plusWord2# l_0 c 393 !h_1 = Exts.plusWord# h_0 d 394 in (# Limb l_1, Limb h_1 #) 395 where 396 -- wide wrapping addition 397 wadd_w# :: (# Word#, Word# #) -> Word# -> (# Word#, Word# #) 398 wadd_w# (# x_lo, x_hi #) y_lo = 399 let !(# c0, s0 #) = Exts.plusWord2# x_lo y_lo 400 !(# _, s1 #) = Exts.plusWord2# x_hi c0 401 in (# s0, s1 #) 402 {-# INLINE wadd_w# #-} 403 {-# INLINE mac# #-} 404