------------------------------------------------------------------------
-- |
-- Module           : What4.Domains.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.Domains.Arithmetic
  ( ctz
  , clz
  , intLog2
  , isPow2Integer
  , bitsBelow
  , rotateLeft
  , rotateRight
  ) where

import Data.Bits (Bits(..), xor, shiftL, shiftR)

import Data.Parameterized.NatRepr

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

-- | /O(w)/. Count trailing zeros, capped at the width.
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

-- | /O(w)/. Count leading zeros, capped at the width.
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

-- | /O(w)/. @intLog2 n@ for @n >= 1@: floor of base-2 logarithm. Undefined
-- for @n <= 0@. On GHC 9.0+ this delegates to a primop in @ghc-bignum@
-- (constant-time per limb); on earlier GHCs it uses a shift loop.
intLog2 :: Integer -> Int
intLog2 :: Integer -> Int
intLog2 = Integer -> Int
intLog2Opt
{-# INLINE intLog2 #-}

-- | /O(w)/. Test whether @n@ is a positive power of two. On GHC 9.0+ this
-- uses the @integerIsPowerOf2#@ primop; on earlier GHCs it uses
-- @n .&. (n - 1) == 0@.
isPow2Integer :: Integer -> Bool
isPow2Integer :: Integer -> Bool
isPow2Integer = Integer -> Bool
isPow2IntegerOpt
{-# INLINE isPow2Integer #-}

-- | /O(w)/. @bitsBelow n@ returns the smallest mask of the form @2^k - 1@
-- that is at least @n@. That is, @2^(floor(log2 n) + 1) - 1@ for @n > 0@,
-- or @0@ for @n <= 0@. Every value in @[0..n]@ has all its set bits within
-- this mask.
bitsBelow :: Integer -> Integer
bitsBelow :: Integer -> Integer
bitsBelow Integer
n
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0    = Integer
0
  | Bool
otherwise = Int -> Integer
forall a. Bits a => Int -> a
bit (Integer -> Int
intLog2 Integer
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
{-# INLINE bitsBelow #-}

-- | /O(w)/. Rotate a @w@-bit value right by @n@ positions (mod @w@).
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)

-- | /O(w)/. Rotate a @w@-bit value left by @n@ positions (mod @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)