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