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 (8297B)


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