    ------------------------------------------------------------------------
    -- |
    -- 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 #-}
    {-# LANGUAGE CPP #-}
    module What4.Utils.Arithmetic
      ( -- * Arithmetic utilities
        isPow2
      , lg
      , lgCeil
      , 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

    -- | 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 floor of log base 2.
    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)

    -- | Returns ceil of log base 2.
    --   We define @lgCeil 0 = 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."

    -- | Count trailing zeros
    ctz :: NatRepr w -> Integer -> Integer
    ctz :: forall (w :: Nat). NatRepr w -> Integer -> Integer
    ctz 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
    iInteger -> Integer -> Integer
    forall a. Num a => a -> a -> a
    +Integer
    1)
       | Bool
    otherwise = Integer
    i

    -- | Count leading zeros
    clz :: NatRepr w -> Integer -> Integer
    clz :: forall (w :: Nat). NatRepr w -> Integer -> Integer
    clz 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
    iInteger -> Integer -> Integer
    forall a. Num a => a -> a -> a
    +Integer
    1)
       | Bool
    otherwise = Integer
    i

    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)