------------------------------------------------------------------------
-- |
-- Module           : What4.Utils.Arithmetic
-- Description      : Utility functions for computing arithmetic
-- Copyright        : (c) Galois, Inc 2015-2020
-- License          : BSD3
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
-- Stability        : provisional
------------------------------------------------------------------------
{-# LANGUAGE BangPatterns #-}
module What4.Utils.Arithmetic
  ( -- * Arithmetic utilities
    isPow2
  , isPow2Integer
  , lg
  , intLog2
  , lgCeil
  , intLogCeil
  , nextMultiple
  , nextPow2Multiple
  , tryIntSqrt
  , tryRationalSqrt
  , roundAway
  , ctz
  , clz
  , rotateLeft
  , rotateRight
  ) where

import Control.Exception (assert)
import Data.Bits (Bits(..))
import Data.Ratio

import Data.Parameterized.NatRepr

import What4.Domains.Arithmetic.Internal
  ( ctzOpt, clzOpt, intLog2Opt, isPow2IntegerOpt )

-- | Returns true if number is a power of two.
isPow2 :: (Bits a, Num a) => a -> Bool
isPow2 :: forall a. (Bits a, Num a) => a -> Bool
isPow2 a
x = a
x a -> a -> a
forall a. Bits a => a -> a -> a
.&. (a
xa -> a -> a
forall a. Num a => a -> a -> a
-a
1) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0

-- | Returns true if Integer is a power of two. On GHC 9.0+ this uses a fast
-- primop from @ghc-bignum@; on earlier GHCs it falls back to 'isPow2'.
isPow2Integer :: Integer -> Bool
isPow2Integer :: Integer -> Bool
isPow2Integer = Integer -> Bool
isPow2IntegerOpt
{-# INLINE isPow2Integer #-}

-- | Returns floor of log base 2. Polymorphic over bit-like types.
--
-- Note: For @Integer@ specifically, prefer 'intLog2' which uses fast primops
-- on GHC 9.0+.
lg :: (Bits a, Num a, Ord a) => a -> Int
lg :: forall a. (Bits a, Num a, Ord a) => a -> Int
lg a
i0 | a
i0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = Int -> a -> Int
forall {t} {t}. (Num t, Num t, Bits t) => t -> t -> t
go Int
0 (a
i0 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
      | Bool
otherwise = [Char] -> Int
forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"lg given number that is not positive."
  where go :: t -> t -> t
go t
r t
0 = t
r
        go t
r t
n = t -> t -> t
go (t
rt -> t -> t
forall a. Num a => a -> a -> a
+t
1) (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)

-- | @intLog2 n@ for @n >= 1@: floor of base-2 logarithm. Undefined for
-- @n <= 0@. On GHC 9.0+ this delegates to a fast primop in @ghc-bignum@;
-- on earlier GHCs it falls back to 'lg'.
intLog2 :: Integer -> Int
intLog2 :: Integer -> Int
intLog2 = Integer -> Int
intLog2Opt
{-# INLINE intLog2 #-}

-- | Returns ceil of log base 2. Polymorphic over bit-like types.
--   We define @lgCeil 0 = 0@ and @lgCeil 1 = 0@.
--
-- Note: For @Integer@ specifically, prefer 'intLogCeil' which uses fast primops
-- on GHC 9.0+.
lgCeil :: (Bits a, Num a, Ord a) => a -> Int
lgCeil :: forall a. (Bits a, Num a, Ord a) => a -> Int
lgCeil a
0 = Int
0
lgCeil a
1 = Int
0
lgCeil a
i | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
1 = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
forall a. (Bits a, Num a, Ord a) => a -> Int
lg (a
ia -> a -> a
forall a. Num a => a -> a -> a
-a
1)
         | Bool
otherwise = [Char] -> Int
forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"lgCeil given number that is not positive."

-- | @intLogCeil n@ for @n >= 0@: ceiling of base-2 logarithm. We define
-- @intLogCeil 0 = 0@ and @intLogCeil 1 = 0@. On GHC 9.0+ this uses fast primops
-- from @ghc-bignum@; on earlier GHCs it falls back to 'lgCeil'.
intLogCeil :: Integer -> Int
intLogCeil :: Integer -> Int
intLogCeil Integer
0 = Int
0
intLogCeil Integer
1 = Int
0
intLogCeil Integer
i | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
intLog2 (Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)
             | Bool
otherwise = [Char] -> Int
forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"intLogCeil given number that is not positive."
{-# INLINE intLogCeil #-}

-- | Count trailing zeros
ctz :: NatRepr w -> Integer -> Integer
ctz :: forall (w :: Nat). NatRepr w -> Integer -> Integer
ctz = NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
ctzOpt

-- | Count leading zeros
clz :: NatRepr w -> Integer -> Integer
clz :: forall (w :: Nat). NatRepr w -> Integer -> Integer
clz = NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
clzOpt

rotateRight ::
  NatRepr w {- ^ width -} ->
  Integer {- ^ value to rotate -} ->
  Integer {- ^ amount to rotate -} ->
  Integer
rotateRight :: forall (w :: Nat). NatRepr w -> Integer -> Integer -> Integer
rotateRight NatRepr w
w Integer
x Integer
n = Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
xor (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
x' Int
n') (NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
toUnsigned NatRepr w
w (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftL 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
- Int
n')))
 where
 x' :: Integer
x' = NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
toUnsigned NatRepr w
w Integer
x
 n' :: Int
n' = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` NatRepr w -> Integer
forall (n :: Nat). NatRepr n -> Integer
intValue NatRepr w
w)

rotateLeft ::
  NatRepr w {- ^ width -} ->
  Integer {- ^ value to rotate -} ->
  Integer {- ^ amount to rotate -} ->
  Integer
rotateLeft :: forall (w :: Nat). NatRepr w -> Integer -> Integer -> Integer
rotateLeft NatRepr w
w Integer
x Integer
n = Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
xor (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR 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
- Int
n')) (NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
toUnsigned NatRepr w
w (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftL Integer
x' Int
n'))
 where
 x' :: Integer
x' = NatRepr w -> Integer -> Integer
forall (w :: Nat). NatRepr w -> Integer -> Integer
toUnsigned NatRepr w
w Integer
x
 n' :: Int
n' = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` NatRepr w -> Integer
forall (n :: Nat). NatRepr n -> Integer
intValue NatRepr w
w)


-- | @nextMultiple x y@ computes the next multiple m of x s.t. m >= y.  E.g.,
-- nextMultiple 4 8 = 8 since 8 is a multiple of 8; nextMultiple 4 7 = 8;
-- nextMultiple 8 6 = 8.
nextMultiple :: Integral a => a -> a -> a
nextMultiple :: forall a. Integral a => a -> a -> a
nextMultiple a
x a
y = ((a
y a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
x) a -> a -> a
forall a. Num a => a -> a -> a
* a
x

-- | @nextPow2Multiple x n@ returns the smallest multiple of @2^n@
-- not less than @x@.
nextPow2Multiple :: (Bits a, Integral a) => a -> Int -> a
nextPow2Multiple :: forall a. (Bits a, Integral a) => a -> Int -> a
nextPow2Multiple a
x Int
n | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = ((a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
2a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^Int
n a -> a -> a
forall a. Num a => a -> a -> a
-a
1) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
n) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
n
                     | Bool
otherwise = [Char] -> a
forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"nextPow2Multiple given negative value."

------------------------------------------------------------------------
-- Sqrt operators.

-- | This returns the sqrt of an integer if it is well-defined.
tryIntSqrt :: Integer -> Maybe Integer
tryIntSqrt :: Integer -> Maybe Integer
tryIntSqrt Integer
0 = Integer -> Maybe Integer
forall a. a -> Maybe a
forall (m :: Type -> Type) a. Monad m => a -> m a
return Integer
0
tryIntSqrt Integer
1 = Integer -> Maybe Integer
forall a. a -> Maybe a
forall (m :: Type -> Type) a. Monad m => a -> m a
return Integer
1
tryIntSqrt Integer
2 = Maybe Integer
forall a. Maybe a
Nothing
tryIntSqrt Integer
3 = Maybe Integer
forall a. Maybe a
Nothing
tryIntSqrt Integer
n = Bool -> Maybe Integer -> Maybe Integer
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
4) (Maybe Integer -> Maybe Integer) -> Maybe Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Maybe Integer
go (Integer
n Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
  where go :: Integer -> Maybe Integer
go Integer
x | Integer
x2 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n  = Maybe Integer
forall a. Maybe a
Nothing   -- Guess is below sqrt, so we quit.
             | Integer
x2 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n = Integer -> Maybe Integer
forall a. a -> Maybe a
forall (m :: Type -> Type) a. Monad m => a -> m a
return Integer
x' -- We have found sqrt
             | Bool
True    = Integer -> Maybe Integer
go Integer
x'     -- Guess is still too large, so try again.
          where -- Next guess is floor(avg(x, n/x))
                x' :: Integer
x' = (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
x) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2
                x2 :: Integer
x2 = Integer
x' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x'

-- | Return the rational sqrt of a
tryRationalSqrt :: Rational -> Maybe Rational
tryRationalSqrt :: Rational -> Maybe Rational
tryRationalSqrt Rational
r = do
  Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
(%) (Integer -> Integer -> Rational)
-> Maybe Integer -> Maybe (Integer -> Rational)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Integer -> Maybe Integer
tryIntSqrt (Rational -> Integer
forall a. Ratio a -> a
numerator   Rational
r)
      Maybe (Integer -> Rational) -> Maybe Integer -> Maybe Rational
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Integer -> Maybe Integer
tryIntSqrt (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r)

------------------------------------------------------------------------
-- Conversion

-- | Evaluate a real to an integer with rounding away from zero.
roundAway :: (RealFrac a) => a -> Integer
roundAway :: forall a. RealFrac a => a -> Integer
roundAway a
r = a -> Integer
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (a
r a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Num a => a -> a
signum a
r a -> a -> a
forall a. Num a => a -> a -> a
* a
0.5)