{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedSums #-}

-- | Internal module exposing both optimized and reference implementations
-- for property testing. Items in this module should /not/ be considered part
-- of What4's API; they are exported only for the sake of the test suite.
module What4.Domains.Arithmetic.Internal
  ( -- * Reference implementations (always available)
    ctzRef
  , clzRef
  , intLog2Ref
  , isPow2IntegerRef
    -- * Optimized implementations (GHC 9.0+ only)
  , ctzOpt
  , clzOpt
  , intLog2Opt
  , isPow2IntegerOpt
  ) where

import Data.Bits (Bits(..), testBit, shiftR)

import Data.Parameterized.NatRepr

#if MIN_VERSION_base(4,15,0)
import qualified GHC.Num.Integer as Integer
import qualified GHC.Num.BigNat as BigNat
import GHC.Exts (Word(..), ctz#, int2Word#)
#endif

------------------------------------------------------------------------
-- Reference implementations (naive loop-based)

-- | Reference implementation: Count trailing zeros using bit testing loop
ctzRef :: NatRepr w -> Integer -> Integer
ctzRef :: forall (w :: Nat). NatRepr w -> Integer -> Integer
ctzRef NatRepr w
w Integer
x = Integer -> Integer
go Integer
0
 where
 go :: Integer -> Integer
go Integer
i
   | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w) Bool -> Bool -> Bool
&& Integer -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Integer
x (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i) Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
False = Integer -> Integer
go (Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
   | Bool
otherwise = Integer
i
{-# INLINABLE ctzRef #-}

-- | Reference implementation: Count leading zeros using bit testing loop
clzRef :: NatRepr w -> Integer -> Integer
clzRef :: forall (w :: Nat). NatRepr w -> Integer -> Integer
clzRef NatRepr w
w Integer
x = Integer -> Integer
go Integer
0
 where
 go :: Integer -> Integer
go Integer
i
   | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w) Bool -> Bool -> Bool
&& Integer -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Integer
x (NatRepr w -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr w
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
False = Integer -> Integer
go (Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
   | Bool
otherwise = Integer
i
{-# INLINABLE clzRef #-}

-- | Reference implementation: Floor of log base 2 using shift loop
intLog2Ref :: Integer -> Int
intLog2Ref :: Integer -> Int
intLog2Ref = Int -> Integer -> Int
forall {t} {t}. (Ord t, Num t, Num t, Bits t) => t -> t -> t
go Int
0
  where
  go :: t -> t -> t
go !t
k t
m
    | t
m t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
1    = t
k
    | Bool
otherwise = t -> t -> t
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) (t
m t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
{-# INLINABLE intLog2Ref #-}

-- | Reference implementation: Check if Integer is a positive power of two.
isPow2IntegerRef :: Integer -> Bool
isPow2IntegerRef :: Integer -> Bool
isPow2IntegerRef Integer
x = Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 Bool -> Bool -> Bool
&& Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
{-# INLINE isPow2IntegerRef #-}

------------------------------------------------------------------------
-- Optimized implementations (GHC 9.0+ primops)

-- | Optimized implementation: Count trailing zeros using ghc-bignum primops
ctzOpt :: NatRepr w -> Integer -> Integer
#if MIN_VERSION_base(4,15,0)
ctzOpt :: forall (w :: Nat). NatRepr w -> Integer -> Integer
ctzOpt NatRepr w
w Integer
x
  | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)
  | Bool
otherwise =
      case Integer
x of
        Integer.IS Int#
i# -> Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min (Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)) (Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Integer) -> Word -> Integer
forall a b. (a -> b) -> a -> b
$ Word# -> Word
W# (Word# -> Word#
ctz# (Int# -> Word#
int2Word# Int#
i#)))
        Integer.IN ByteArray#
bn -> Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min (Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)) (Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Integer) -> Word -> Integer
forall a b. (a -> b) -> a -> b
$ ByteArray# -> Word
BigNat.bigNatCtz ByteArray#
bn)
        Integer.IP ByteArray#
bn -> Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min (Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)) (Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Integer) -> Word -> Integer
forall a b. (a -> b) -> a -> b
$ ByteArray# -> Word
BigNat.bigNatCtz ByteArray#
bn)
#else
ctzOpt = ctzRef
#endif
{-# INLINE ctzOpt #-}

-- | Optimized implementation: Count leading zeros using integerLog2 primop
clzOpt :: NatRepr w -> Integer -> Integer
#if MIN_VERSION_base(4,15,0)
clzOpt :: forall (w :: Nat). NatRepr w -> Integer -> Integer
clzOpt NatRepr w
w Integer
x
  | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)
  | Bool
otherwise =
      -- Mask to width-w value to handle negative numbers and values outside range
      let width :: Integer
width = Nat -> Integer
forall a. Integral a => a -> Integer
toInteger (NatRepr w -> Nat
forall (n :: Nat). NatRepr n -> Nat
natValue NatRepr w
w)
          mask :: Integer
mask = (Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
width) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
          x' :: Integer
x' = Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask
      in if Integer
x' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
         then Integer
width
         else let highBit :: Integer
highBit = Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Word
Integer.integerLog2 Integer
x')
              in if Integer
highBit Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
width
                 then Integer
0
                 else Integer
width Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
highBit
#else
clzOpt = clzRef
#endif
{-# INLINE clzOpt #-}

-- | Optimized implementation: Floor of log base 2 using integerLog2 primop
intLog2Opt :: Integer -> Int
#if MIN_VERSION_base(4,15,0)
intLog2Opt :: Integer -> Int
intLog2Opt Integer
n = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Word
Integer.integerLog2 Integer
n)
#else
intLog2Opt = intLog2Ref
#endif
{-# INLINE intLog2Opt #-}

-- | Optimized implementation: Check if Integer is power of two using primops
isPow2IntegerOpt :: Integer -> Bool
#if MIN_VERSION_base(4,15,0)
isPow2IntegerOpt :: Integer -> Bool
isPow2IntegerOpt Integer
x = case Integer -> (# (# #) | Word# #)
Integer.integerIsPowerOf2# Integer
x of
  (# (# #)
_ | #) -> Bool
False
  (# | Word#
_ #) -> Bool
True
#else
isPow2IntegerOpt = isPow2IntegerRef
#endif
{-# INLINE isPow2IntegerOpt #-}