-- | Bit manipulation

module Feldspar.Core.Functions.Bits where

import qualified Data.Bits as B
import Data.Int
import Data.Word

import Feldspar.Range
import Feldspar.Core.Types
import Feldspar.Core.Representation
import Feldspar.Core.Constructs

infixl 5 <<,>>
infixl 4 ⊕

-- | Redefinition of the standard 'B.Bits' class for Feldspar
class (B.Bits a, Type a, FullProp (Size a)) => Bits a
  where
  -- Logical operations
  (.&.)         :: Data a -> Data a -> Data a
  (.&.)         =  optAnd fullProp
  (.|.)         :: Data a -> Data a -> Data a
  (.|.)         =  optOr fullProp
  xor           :: Data a -> Data a -> Data a
  xor           =  optXor fullProp
  complement    :: Data a -> Data a
  complement    =  function1 "complement" fullProp B.complement

  -- Operations on individual bits
  bit           :: Data Index -> Data a
  bit           =  function1 "bit" fullProp (B.bit . fromIntegral)
  setBit        :: Data a -> Data Index -> Data a
  setBit        =  function2 "setBit" fullProp (liftIntWord B.setBit)
  clearBit      :: Data a -> Data Index -> Data a
  clearBit      =  function2 "clearBit" fullProp (liftIntWord B.clearBit)
  complementBit :: Data a -> Data Index -> Data a
  complementBit =  function2 "complementBit" fullProp (liftIntWord B.complementBit)
  testBit       :: Data a -> Data Index -> Data Bool
  testBit       =  function2 "testBit" fullProp (liftIntWord B.testBit)

  -- Moving bits around
  shiftLU       :: Data a -> Data Index -> Data a
  shiftLU       =  optZero $ function2 "shiftL" fullProp (liftIntWord B.shiftL)
  shiftRU       :: Data a -> Data Index -> Data a
  shiftRU       =  optZero $ function2 "shiftR" fullProp cShiftRU
  shiftL        :: Data a -> Data DefaultInt -> Data a
  shiftL        =  optZero $ function2 "shiftL" fullProp (liftInt B.shiftL)
  shiftR        :: Data a -> Data DefaultInt -> Data a
  shiftR        =  optZero $ function2 "shiftR" fullProp (liftInt B.shiftR)
  rotateLU      :: Data a -> Data Index -> Data a
  rotateLU      =  optZero $ function2 "rotateL" fullProp (liftIntWord B.rotateL )
  rotateRU      :: Data a -> Data Index -> Data a
  rotateRU      =  optZero $ function2 "rotateR" fullProp (liftIntWord B.rotateR )
  rotateL       :: Data a -> Data DefaultInt -> Data a
  rotateL       =  optZero $ function2 "rotateL" fullProp (liftInt B.rotateL)
  rotateR       :: Data a -> Data DefaultInt -> Data a
  rotateR       =  optZero $ function2 "rotateR" fullProp (liftInt B.rotateR)
  reverseBits   :: Data a -> Data a
  reverseBits   =  function1 "reverseBits" fullProp evalReverseBits

  -- Bulk bit operations
  -- | Returns the number of leading zeroes for unsigned types.
  -- For signed types it returns the number of unnecessary sign bits
  bitScan       :: Data a -> Data Index
  bitScan       =  function1 "bitScan" fullProp (fromIntegral . evalBitScan)
  bitCount      :: Data a -> Data Index
  bitCount      =  function1 "bitCount" fullProp (fromIntegral . evalBitCount)

  -- Queries about the type
  bitSize       :: Data a -> Data Index
  bitSize       =  function1 "bitSize" (\_ -> naturalRange) (fromIntegral . B.bitSize)
  isSigned      :: Data a -> Data Bool
  isSigned      =  function1 "isSigned" fullProp B.isSigned

-- TODO Some range propagation could be improved.
--      bitSize could have (range 0 32) instead of naturalRange.

liftIntWord :: (a -> Int -> b) -> (a -> DefaultWord -> b)
liftIntWord f x = f x . fromIntegral

liftInt :: (a -> Int -> b) -> (a -> DefaultInt -> b)
liftInt f x = f x . fromIntegral

(⊕)  :: Bits a => Data a -> Data a -> Data a
(⊕)  =  xor
(<<) :: Bits a => Data a -> Data Index -> Data a
(<<) =  shiftLU
(>>) :: Bits a => Data a -> Data Index -> Data a
(>>) =  shiftRU

optAnd :: (Bits a) =>
          (Size a -> Size a -> Size a)
       -> Data a -> Data a -> Data a
optAnd rangeProp x y =
    case (viewLiteral x, viewLiteral y) of
      (Just 0, _) -> value 0
      (_, Just 0) -> value 0
      (Just x, _) | isAllOnes x -> y
      (_, Just y) | isAllOnes y -> x
      _ -> function2 "(.&.)" rangeProp (B..&.) x y

optOr :: (Bits a) =>
         (Size a -> Size a -> Size a)
      -> Data a -> Data a -> Data a
optOr rangeProp x y =
    case (viewLiteral x, viewLiteral y) of
      (Just 0, _) -> y
      (_, Just 0) -> x
      (Just x, _) | isAllOnes x -> value allOnes
      (_, Just y) | isAllOnes y -> value allOnes
      _ -> function2 "(.|.)" rangeProp (B..|.) x y

optXor :: (Bits a) =>
          (Size a -> Size a -> Size a)
       -> Data a -> Data a -> Data a
optXor rangeProp x y =
    case (viewLiteral x, viewLiteral y) of
      (Just 0, _) -> y
      (_, Just 0) -> x
      (Just x, _) | isAllOnes x -> complement y
      (_, Just y) | isAllOnes y -> complement x
      _ -> function2 "xor" rangeProp B.xor x y

isAllOnes :: B.Bits a => a -> Bool
isAllOnes x = x Prelude.== B.complement 0

allOnes :: B.Bits a => a
allOnes = B.complement 0

optZero :: (Type n, Num n) => (a -> Data n -> a) -> a -> Data n -> a
optZero f x y = case viewLiteral y of
                  Just 0 -> x
                  _      -> f x y

evalBitScan :: B.Bits b => b -> Word
evalBitScan b =
   if B.isSigned b
   then scanLoop b (B.testBit b (B.bitSize b - 1)) (B.bitSize b - 2) 0
         else scanLoop b False (B.bitSize b - 1) 0
  where
    scanLoop b bit i n | i Prelude.< 0                = n
    scanLoop b bit i n | B.testBit b i Prelude./= bit = n
    scanLoop b bit i n | otherwise                    = scanLoop b bit (i-1) (n+1)

evalBitCount :: B.Bits b => b -> Word
evalBitCount b = loop b (B.bitSize b - 1) 0
  where
    loop b i n | i Prelude.< 0 = n
    loop b i n | B.testBit b i = loop b (i-1) (n+1)
    loop b i n | otherwise     = loop b (i-1) n

evalReverseBits :: B.Bits b => b -> b
evalReverseBits b = revLoop b 0 (0 `asTypeOf` b)
  where
    bitSize = B.bitSize b
    revLoop b i n | i Prelude.>= bitSize  = n
    revLoop b i n | B.testBit b i = revLoop b (i+1) (B.setBit n (bitSize - i - 1))
    revLoop b i n | otherwise     = revLoop b (i+1) n

{- TODO
   This is a hack until we have proper size information for DefaultWord.
   We do not want the Range module to have to know about DefaultWord so
   we patch the types for this range propagation function here.
-}
propRangeShiftLU r1 r2
    = rangeShiftLU r1 (mapMonotonic (\ (DefaultWord a) -> a) r2)
propRangeShiftRU r1 r2
    = rangeShiftRU r1 (mapMonotonic (\ (DefaultWord a) -> a) r2)
cShiftRU v (DefaultWord i) = correctShiftRU v i

{- TODO
   The reason we have to provide the range propagation functions
   in these instances is that if we would try to do it in the default
   methods of the class we would get a superclass constraint 'Size a ~ Range a'.
   GHC 6.12 doesn't support this.
-}
instance Bits Word8 where
  xor     = optXor rangeXor
  shiftLU = optZero $ function2 "shiftL" propRangeShiftLU (liftIntWord B.shiftL)
  shiftRU = optZero $ function2 "shiftR" propRangeShiftRU cShiftRU
instance Bits Int8 where
  xor = optXor rangeXor
instance Bits Word16 where
  xor = optXor rangeXor
  shiftLU = optZero $ function2 "shiftL" propRangeShiftLU (liftIntWord B.shiftL)
  shiftRU = optZero $ function2 "shiftR" propRangeShiftRU cShiftRU
instance Bits Int16 where
  xor = optXor rangeXor
instance Bits Word32 where
  xor = optXor rangeXor
  shiftLU = optZero $ function2 "shiftL" propRangeShiftLU (liftIntWord B.shiftL)
  shiftRU = optZero $ function2 "shiftR" propRangeShiftRU cShiftRU
instance Bits Int32 where
  xor = optXor rangeXor
instance Bits DefaultWord where
  xor = optXor rangeXor
  shiftLU = optZero $ function2 "shiftL" propRangeShiftLU (liftIntWord B.shiftL)
  shiftRU = optZero $ function2 "shiftR" propRangeShiftRU cShiftRU
instance Bits DefaultInt where
  xor = optXor rangeXor