{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE UnboxedSums #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedNewtypes #-}

-- |
-- Module: Data.Word.Wide
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Wide words, consisting of two 'Limb's.

module Data.Word.Wide (
  -- * Wide Words
    Wide(..)

  -- * Construction, Conversion
  , wide
  , to_vartime
  , from_vartime

  -- * Constant-time selection
  , select
  , select#

  -- * Bit Manipulation
  , or
  , or#
  , and
  , and#
  , xor
  , xor#
  , not
  , not#

  -- * Comparison
  , eq
  , eq_vartime

  -- * Arithmetic
  , add
  , add_o
  , sub
  , mul
  , neg

  -- * Unboxed Arithmetic
  , add_o#
  , add_w#
  , sub_b#
  , sub_w#
  , mul_w#
  , neg#
  ) where

import Control.DeepSeq
import Data.Bits ((.|.), (.&.), (.<<.), (.>>.))
import qualified Data.Bits as B
import qualified Data.Choice as C
import Data.Word.Limb (Limb(..))
import qualified Data.Word.Limb as L
import Prelude hiding (div, mod, or, and, not, quot, rem, recip)

-- utilities ------------------------------------------------------------------

fi :: (Integral a, Num b) => a -> b
fi :: forall a b. (Integral a, Num b) => a -> b
fi = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE fi #-}

-- wide words -----------------------------------------------------------------

type Limb2 = (# Limb, Limb #)

pattern L2 :: L.Word# -> L.Word# -> Limb2
pattern $mL2 :: forall {r}. Limb2 -> (Word# -> Word# -> r) -> ((# #) -> r) -> r
$bL2 :: Word# -> Word# -> Limb2
L2 w0 w1 = (# Limb w0, Limb w1 #)
{-# COMPLETE L2 #-}

-- | Little-endian wide words.
data Wide = Wide !Limb2

instance Show Wide where
  show :: Wide -> String
show = Integer -> String
forall a. Show a => a -> String
show (Integer -> String) -> (Wide -> Integer) -> Wide -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Wide -> Integer
from_vartime

-- | Note that 'fromInteger' necessarily runs in variable time due
--   to conversion from the variable-size, potentially heap-allocated
--   'Integer' type.
instance Num Wide where
  + :: Wide -> Wide -> Wide
(+) = Wide -> Wide -> Wide
add
  (-) = Wide -> Wide -> Wide
sub
  * :: Wide -> Wide -> Wide
(*) = Wide -> Wide -> Wide
mul
  abs :: Wide -> Wide
abs = Wide -> Wide
forall a. a -> a
id
  fromInteger :: Integer -> Wide
fromInteger = Integer -> Wide
to_vartime
  negate :: Wide -> Wide
negate = Wide -> Wide
neg
  signum :: Wide -> Wide
signum (Wide (# Limb
l0, Limb
l1 #)) =
    let !(Limb Word#
l) = Limb
l0 Limb -> Limb -> Limb
`L.or#` Limb
l1
        !n :: Choice
n = Word# -> Choice
C.from_word_nonzero# Word#
l
        !b :: Word#
b = Choice -> Word#
C.to_word# Choice
n
    in  Limb2 -> Wide
Wide (Word# -> Word# -> Limb2
L2 Word#
b Word#
0##)

instance NFData Wide where
  rnf :: Wide -> ()
rnf (Wide Limb2
a) = case Limb2
a of (# Limb
_, Limb
_ #) -> ()

-- construction / conversion --------------------------------------------------

-- | Construct a 'Wide' word from low and high 'Word's.
wide :: Word -> Word -> Wide
wide :: Word -> Word -> Wide
wide (L.W# Word#
l) (L.W# Word#
h) = Limb2 -> Wide
Wide (# Word# -> Limb
Limb Word#
l, Word# -> Limb
Limb Word#
h #)
{-# INLINE wide #-}

-- | Convert an 'Integer' to a 'Wide' word in variable time.
--
--   >>> to_vartime 1
--   1
to_vartime :: Integer -> Wide
to_vartime :: Integer -> Wide
to_vartime Integer
n =
  let !size :: Int
size = Word -> Int
forall b. FiniteBits b => b -> Int
B.finiteBitSize (Word
0 :: Word)
      !mask :: Integer
mask = Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fi (Word
forall a. Bounded a => a
maxBound :: Word) :: Integer
      !(L.W# Word#
w0) = Integer -> Word
forall a b. (Integral a, Num b) => a -> b
fi (Integer
n Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask)
      !(L.W# Word#
w1) = Integer -> Word
forall a b. (Integral a, Num b) => a -> b
fi ((Integer
n Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
.>>. Int
size) Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask)
  in  Limb2 -> Wide
Wide (# Word# -> Limb
Limb Word#
w0, Word# -> Limb
Limb Word#
w1 #)
{-# INLINABLE to_vartime #-}

-- | Convert a 'Wide' word to an 'Integer' in variable time.
--
--   >>> from_vartime 1
--   1
from_vartime :: Wide -> Integer
from_vartime :: Wide -> Integer
from_vartime (Wide (# Limb Word#
a, Limb Word#
b #)) =
      Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fi (Word# -> Word
L.W# Word#
b) Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
.<<. (Word -> Int
forall b. FiniteBits b => b -> Int
B.finiteBitSize (Word
0 :: Word))
  Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.|. Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fi (Word# -> Word
L.W# Word#
a)
{-# INLINABLE from_vartime #-}

-- comparison -----------------------------------------------------------------

-- | Compare 'Wide' words for equality in constant time.
--
--   >>> import qualified Data.Chocie as C
--   >>> C.decide (eq 1 1)
--   True
eq :: Wide -> Wide -> C.Choice
eq :: Wide -> Wide -> Choice
eq (Wide (# Limb Word#
a0, Limb Word#
a1 #)) (Wide (# Limb Word#
b0, Limb Word#
b1 #)) =
  Limb2 -> Limb2 -> Choice
C.eq_wide# (# Word#
a0, Word#
a1 #) (# Word#
b0, Word#
b1 #)
{-# INLINABLE eq #-}

-- | Compare 'Wide' words for equality in variable time.
--
--   >>> eq_vartime 1 1
--   True
eq_vartime :: Wide -> Wide -> Bool
eq_vartime :: Wide -> Wide -> Bool
eq_vartime (Wide (# Limb
a0, Limb
b0 #)) (Wide (# Limb
a1, Limb
b1 #)) =
  Limb -> Limb -> Bool
L.eq_vartime# Limb
a0 Limb
a1 Bool -> Bool -> Bool
&& Limb -> Limb -> Bool
L.eq_vartime# Limb
b0 Limb
b1
{-# INLINABLE eq_vartime #-}

-- constant-time selection-----------------------------------------------------

-- | Return a if c is truthy, otherwise return b.
--
--   >>> import qualified Data.Choice as C
--   >>> select 0 1 (C.true# ())
--   1
select
  :: Wide     -- ^ a
  -> Wide     -- ^ b
  -> C.Choice -- ^ c
  -> Wide     -- ^ result
select :: Wide -> Wide -> Choice -> Wide
select (Wide Limb2
a) (Wide Limb2
b) Choice
c = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Choice -> Limb2
select# Limb2
a Limb2
b Choice
c)
{-# INLINABLE select #-}

select#
  :: Limb2    -- ^ a
  -> Limb2    -- ^ b
  -> C.Choice -- ^ c
  -> Limb2    -- ^ result
select# :: Limb2 -> Limb2 -> Choice -> Limb2
select# (L2 Word#
a0 Word#
a1) (L2 Word#
b0 Word#
b1) Choice
c =
  let !(# Word#
w0, Word#
w1 #) = Limb2 -> Limb2 -> Choice -> Limb2
C.select_wide# (# Word#
a0, Word#
a1 #) (# Word#
b0, Word#
b1 #) Choice
c
  in  Word# -> Word# -> Limb2
L2 Word#
w0 Word#
w1
{-# INLINE select# #-}

-- bits -----------------------------------------------------------------------

or# :: Limb2 -> Limb2 -> Limb2
or# :: Limb2 -> Limb2 -> Limb2
or# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) = (# Limb -> Limb -> Limb
L.or# Limb
a0 Limb
b0, Limb -> Limb -> Limb
L.or# Limb
a1 Limb
b1 #)
{-# INLINE or# #-}

-- | Logical disjunction on 'Wide' words.
or :: Wide -> Wide -> Wide
or :: Wide -> Wide -> Wide
or (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
or# Limb2
a Limb2
b)
{-# INLINABLE or #-}

and# :: Limb2 -> Limb2 -> Limb2
and# :: Limb2 -> Limb2 -> Limb2
and# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) = (# Limb -> Limb -> Limb
L.and# Limb
a0 Limb
b0, Limb -> Limb -> Limb
L.and# Limb
a1 Limb
b1 #)
{-# INLINE and# #-}

-- | Logical conjunction on 'Wide' words.
and :: Wide -> Wide -> Wide
and :: Wide -> Wide -> Wide
and (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
and# Limb2
a Limb2
b)
{-# INLINABLE and #-}

xor# :: Limb2 -> Limb2 -> Limb2
xor# :: Limb2 -> Limb2 -> Limb2
xor# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) = (# Limb -> Limb -> Limb
L.xor# Limb
a0 Limb
b0, Limb -> Limb -> Limb
L.xor# Limb
a1 Limb
b1 #)
{-# INLINE xor# #-}

-- | Logical exclusive-or on 'Wide' words.
xor :: Wide -> Wide -> Wide
xor :: Wide -> Wide -> Wide
xor (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
xor# Limb2
a Limb2
b)
{-# INLINABLE xor #-}

not# :: Limb2 -> Limb2
not# :: Limb2 -> Limb2
not# (# Limb
a0, Limb
a1 #) = (# Limb -> Limb
L.not# Limb
a0, Limb -> Limb
L.not# Limb
a1 #)
{-# INLINE not# #-}

-- | Logical negation on 'Wide' words.
not :: Wide -> Wide
not :: Wide -> Wide
not (Wide Limb2
w) = Limb2 -> Wide
Wide (Limb2 -> Limb2
not# Limb2
w)
{-# INLINABLE not #-}

-- negation -------------------------------------------------------------------

-- | Wrapping negation on 'Wide' words, producing an additive inverse.
--
--   >>> neg 1
--   340282366920938463463374607431768211455
--   >>> 1 + neg 1
--   >>> 0
neg
  :: Wide -- ^ argument
  -> Wide -- ^ (wrapping) additive inverse
neg :: Wide -> Wide
neg (Wide Limb2
w) = Limb2 -> Wide
Wide (Limb2 -> Limb2
neg# Limb2
w)
{-# INLINABLE neg #-}

neg#
  :: Limb2 -- ^ argument
  -> Limb2 -- ^ (wrapping) additive inverse
neg# :: Limb2 -> Limb2
neg# Limb2
w = Limb2 -> Limb2 -> Limb2
add_w# (Limb2 -> Limb2
not# Limb2
w) (Word# -> Word# -> Limb2
L2 Word#
1## Word#
0##)
{-# INLINE neg# #-}

-- addition, subtraction ------------------------------------------------------

-- | Overflowing addition, computing 'a + b', returning the sum and a
--   carry bit.
add_o#
  :: Limb2              -- ^ augend
  -> Limb2              -- ^ addend
  -> (# Limb2, Limb #)  -- ^ (# sum, carry bit #)
add_o# :: Limb2 -> Limb2 -> (# Limb2, Limb #)
add_o# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) =
  let !(# Limb
s0, Limb
c0 #) = Limb -> Limb -> Limb2
L.add_o# Limb
a0 Limb
b0
      !(# Limb
s1, Limb
c1 #) = Limb -> Limb -> Limb -> Limb2
L.add_c# Limb
a1 Limb
b1 Limb
c0
  in  (# (# Limb
s0, Limb
s1 #), Limb
c1 #)
{-# INLINE add_o# #-}

-- | Overflowing addition on 'Wide' words, computing 'a + b', returning
--   the sum and carry bit.
add_o
  :: Wide         -- ^ augend
  -> Wide         -- ^ addend
  -> (Wide, Word) -- ^ (sum, carry)
add_o :: Wide -> Wide -> (Wide, Word)
add_o (Wide Limb2
a) (Wide Limb2
b) =
  let !(# Limb2
s, Limb Word#
c #) = Limb2 -> Limb2 -> (# Limb2, Limb #)
add_o# Limb2
a Limb2
b
  in  (Limb2 -> Wide
Wide Limb2
s, Word# -> Word
L.W# Word#
c)

-- | Wrapping addition, computing 'a + b'.
add_w#
  :: Limb2 -- ^ augend
  -> Limb2 -- ^ addend
  -> Limb2 -- ^ sum
add_w# :: Limb2 -> Limb2 -> Limb2
add_w# Limb2
a Limb2
b =
  let !(# Limb2
c, Limb
_ #) = Limb2 -> Limb2 -> (# Limb2, Limb #)
add_o# Limb2
a Limb2
b
  in  Limb2
c
{-# INLINE add_w# #-}

-- | Wrapping addition on 'Wide' words, computing 'a + b'.
add :: Wide -> Wide -> Wide
add :: Wide -> Wide -> Wide
add (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
add_w# Limb2
a Limb2
b)

-- | Borrowing subtraction, computing 'a - b' and returning the
--   difference with a borrow mask.
sub_b#
  :: Limb2              -- ^ minuend
  -> Limb2              -- ^ subtrahend
  -> (# Limb2, Limb #) -- ^ (# difference, borrow mask #)
sub_b# :: Limb2 -> Limb2 -> (# Limb2, Limb #)
sub_b# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) =
  let !(# Limb
s0, Limb
c0 #) = Limb -> Limb -> Limb -> Limb2
L.sub_b# Limb
a0 Limb
b0 (Word# -> Limb
Limb Word#
0##)
      !(# Limb
s1, Limb
c1 #) = Limb -> Limb -> Limb -> Limb2
L.sub_b# Limb
a1 Limb
b1 Limb
c0
  in  (# (# Limb
s0, Limb
s1 #), Limb
c1 #)
{-# INLINE sub_b# #-}

-- | Wrapping subtraction, computing 'a - b'.
sub_w#
  :: Limb2 -- ^ minuend
  -> Limb2 -- ^ subtrahend
  -> Limb2 -- ^ difference
sub_w# :: Limb2 -> Limb2 -> Limb2
sub_w# Limb2
a Limb2
b =
  let !(# Limb2
c, Limb
_ #) = Limb2 -> Limb2 -> (# Limb2, Limb #)
sub_b# Limb2
a Limb2
b
  in  Limb2
c
{-# INLINE sub_w# #-}

-- | Wrapping subtraction on 'Wide' words, computing 'a - b'.
sub :: Wide -> Wide -> Wide
sub :: Wide -> Wide -> Wide
sub (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
sub_w# Limb2
a Limb2
b)

-- multiplication -------------------------------------------------------------

-- | Wrapping multiplication, computing 'a b'.
mul_w#
  :: Limb2 -- ^ multiplicand
  -> Limb2 -- ^ multiplier
  -> Limb2 -- ^ product
mul_w# :: Limb2 -> Limb2 -> Limb2
mul_w# (# Limb
a0, Limb
a1 #) (# Limb
b0, Limb
b1 #) =
  let !(# Limb
p0_lo, Limb
p0_hi #) = Limb -> Limb -> Limb2
L.mul_c# Limb
a0 Limb
b0
      !(# Limb
p1_lo, Limb
_ #) = Limb -> Limb -> Limb2
L.mul_c# Limb
a0 Limb
b1
      !(# Limb
p2_lo, Limb
_ #) = Limb -> Limb -> Limb2
L.mul_c# Limb
a1 Limb
b0
      !(# Limb
s0, Limb
_ #) = Limb -> Limb -> Limb2
L.add_o# Limb
p0_hi Limb
p1_lo
      !(# Limb
s1, Limb
_ #) = Limb -> Limb -> Limb2
L.add_o# Limb
s0 Limb
p2_lo
  in  (# Limb
p0_lo, Limb
s1 #)
{-# INLINE mul_w# #-}

-- | Wrapping multiplication on 'Wide' words.
mul :: Wide -> Wide -> Wide
mul :: Wide -> Wide -> Wide
mul (Wide Limb2
a) (Wide Limb2
b) = Limb2 -> Wide
Wide (Limb2 -> Limb2 -> Limb2
mul_w# Limb2
a Limb2
b)