-- copied & adapted from cryptic
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE Rank2Types          #-}
{-# LANGUAGE GADTs               #-}
module Crypto.Math.Bits
    ( FBits(..)
    , FBitsK(..)
    , SizeValid
    , splitHalf
    , append
    , dropBitsOnRight
    , dropBitsOnLeft
    ) where

import Data.Bits
import Data.Proxy
import GHC.Natural
import GHC.TypeLits

-- | Finite Bits
--
-- Sadly Bits is taken by Bits operation
data FBits (n :: Nat) = FBits { FBits n -> Natural
unFBits :: Natural }
    deriving (Int -> FBits n -> ShowS
[FBits n] -> ShowS
FBits n -> String
(Int -> FBits n -> ShowS)
-> (FBits n -> String) -> ([FBits n] -> ShowS) -> Show (FBits n)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: Nat). Int -> FBits n -> ShowS
forall (n :: Nat). [FBits n] -> ShowS
forall (n :: Nat). FBits n -> String
showList :: [FBits n] -> ShowS
$cshowList :: forall (n :: Nat). [FBits n] -> ShowS
show :: FBits n -> String
$cshow :: forall (n :: Nat). FBits n -> String
showsPrec :: Int -> FBits n -> ShowS
$cshowsPrec :: forall (n :: Nat). Int -> FBits n -> ShowS
Show,FBits n -> FBits n -> Bool
(FBits n -> FBits n -> Bool)
-> (FBits n -> FBits n -> Bool) -> Eq (FBits n)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: Nat). FBits n -> FBits n -> Bool
/= :: FBits n -> FBits n -> Bool
$c/= :: forall (n :: Nat). FBits n -> FBits n -> Bool
== :: FBits n -> FBits n -> Bool
$c== :: forall (n :: Nat). FBits n -> FBits n -> Bool
Eq,Eq (FBits n)
Eq (FBits n)
-> (FBits n -> FBits n -> Ordering)
-> (FBits n -> FBits n -> Bool)
-> (FBits n -> FBits n -> Bool)
-> (FBits n -> FBits n -> Bool)
-> (FBits n -> FBits n -> Bool)
-> (FBits n -> FBits n -> FBits n)
-> (FBits n -> FBits n -> FBits n)
-> Ord (FBits n)
FBits n -> FBits n -> Bool
FBits n -> FBits n -> Ordering
FBits n -> FBits n -> FBits n
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall (n :: Nat). Eq (FBits n)
forall (n :: Nat). FBits n -> FBits n -> Bool
forall (n :: Nat). FBits n -> FBits n -> Ordering
forall (n :: Nat). FBits n -> FBits n -> FBits n
min :: FBits n -> FBits n -> FBits n
$cmin :: forall (n :: Nat). FBits n -> FBits n -> FBits n
max :: FBits n -> FBits n -> FBits n
$cmax :: forall (n :: Nat). FBits n -> FBits n -> FBits n
>= :: FBits n -> FBits n -> Bool
$c>= :: forall (n :: Nat). FBits n -> FBits n -> Bool
> :: FBits n -> FBits n -> Bool
$c> :: forall (n :: Nat). FBits n -> FBits n -> Bool
<= :: FBits n -> FBits n -> Bool
$c<= :: forall (n :: Nat). FBits n -> FBits n -> Bool
< :: FBits n -> FBits n -> Bool
$c< :: forall (n :: Nat). FBits n -> FBits n -> Bool
compare :: FBits n -> FBits n -> Ordering
$ccompare :: forall (n :: Nat). FBits n -> FBits n -> Ordering
$cp1Ord :: forall (n :: Nat). Eq (FBits n)
Ord)

data FBitsK = FBitsK (forall n . (KnownNat n, SizeValid n) => FBits n)

type SizeValid n = (KnownNat n, 1 <= n)

toFBits :: SizeValid n => Natural -> FBits n
toFBits :: Natural -> FBits n
toFBits Natural
nat = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
nat FBits n -> FBits n -> FBits n
forall a. Bits a => a -> a -> a
.&. FBits n
forall (n :: Nat). SizeValid n => FBits n
allOne

instance SizeValid n => Enum (FBits n) where
    toEnum :: Int -> FBits n
toEnum Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
&& Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> FBits n -> Natural
forall (n :: Nat). FBits n -> Natural
unFBits FBits n
maxi = String -> FBits n
forall a. HasCallStack => String -> a
error String
"FBits n not within bound"
             | Bool
otherwise                              = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)
      where maxi :: FBits n
maxi = FBits n
forall (n :: Nat). SizeValid n => FBits n
allOne :: FBits n
    fromEnum :: FBits n -> Int
fromEnum (FBits Natural
n) = Natural -> Int
forall a. Enum a => a -> Int
fromEnum Natural
n

instance SizeValid n => Bounded (FBits n) where
    minBound :: FBits n
minBound = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
0
    maxBound :: FBits n
maxBound = FBits n
forall (n :: Nat). SizeValid n => FBits n
allOne

instance SizeValid n => Num (FBits n) where
    fromInteger :: Integer -> FBits n
fromInteger = Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits (Natural -> FBits n) -> (Integer -> Natural) -> Integer -> FBits n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Natural
forall a. Num a => Integer -> a
fromInteger
    + :: FBits n -> FBits n -> FBits n
(+) (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits (Natural
a Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
b)
    (-) (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits (Natural
a Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
b)
    * :: FBits n -> FBits n -> FBits n
(*) (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits (Natural
a Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
* Natural
b)
    abs :: FBits n -> FBits n
abs = FBits n -> FBits n
forall a. a -> a
id
    signum :: FBits n -> FBits n
signum (FBits Natural
a)
        | Natural
a Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0    = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
0
        | Bool
otherwise = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
1

instance SizeValid n => Bits (FBits n) where
    .&. :: FBits n -> FBits n -> FBits n
(.&.) (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Natural -> Natural
forall a. Bits a => a -> a -> a
.&. Natural
b)
    .|. :: FBits n -> FBits n -> FBits n
(.|.) (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Natural -> Natural
forall a. Bits a => a -> a -> a
.|. Natural
b)
    xor :: FBits n -> FBits n -> FBits n
xor (FBits Natural
a) (FBits Natural
b) = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Natural -> Natural
forall a. Bits a => a -> a -> a
`xor` Natural
b)
    shiftR :: FBits n -> Int -> FBits n
shiftR (FBits Natural
a) Int
n = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`shiftR` Int
n)
    shiftL :: FBits n -> Int -> FBits n
shiftL (FBits Natural
a) Int
n = Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits (Natural
a Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`shiftL` Int
n) -- shiftL can overflow here, so explicit safe reconstruction from natural
    rotateL :: FBits n -> Int -> FBits n
rotateL FBits n
a Int
i = ((FBits n
a FBits n -> Int -> FBits n
forall a. Bits a => a -> Int -> a
`shiftL` Int
i) FBits n -> FBits n -> FBits n
forall a. Bits a => a -> a -> a
.|. (FBits n
a FBits n -> Int -> FBits n
forall a. Bits a => a -> Int -> a
`shiftR` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)))
      where n :: Int
n = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)
    rotateR :: FBits n -> Int -> FBits n
rotateR FBits n
a Int
i = ((FBits n
a FBits n -> Int -> FBits n
forall a. Bits a => a -> Int -> a
`shiftR` Int
i) FBits n -> FBits n -> FBits n
forall a. Bits a => a -> a -> a
.|. (FBits n
a FBits n -> Int -> FBits n
forall a. Bits a => a -> Int -> a
`shiftL` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)))
      where n :: Int
n = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)
    zeroBits :: FBits n
zeroBits = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
0
    bit :: Int -> FBits n
bit Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n) = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits Natural
0
        | Bool
otherwise                                            = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
2Natural -> Int -> Natural
forall a b. (Num a, Integral b) => a -> b -> a
^Int
i)
    testBit :: FBits n -> Int -> Bool
testBit (FBits Natural
a) Int
i = Natural -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Natural
a Int
i
    bitSize :: FBits n -> Int
bitSize FBits n
_ = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)
    bitSizeMaybe :: FBits n -> Maybe Int
bitSizeMaybe FBits n
_ = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)
    isSigned :: FBits n -> Bool
isSigned FBits n
_ = Bool
False
    complement :: FBits n -> FBits n
complement FBits n
a = FBits n
forall (n :: Nat). SizeValid n => FBits n
allOne FBits n -> FBits n -> FBits n
forall a. Bits a => a -> a -> a
`xor` FBits n
a
    popCount :: FBits n -> Int
popCount (FBits Natural
a) = Natural -> Int
forall a. Bits a => a -> Int
popCount Natural
a

allOne :: forall n . SizeValid n => FBits n
allOne :: FBits n
allOne = Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
2 Natural -> Integer -> Natural
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1)
  where n :: Integer
n = Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)

splitHalf :: forall m n . (SizeValid n, (n * 2) ~ m) => FBits m -> (FBits n, FBits n)
splitHalf :: FBits m -> (FBits n, FBits n)
splitHalf (FBits Natural
a) = (Natural -> FBits n
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`shiftR` Int
n), Natural -> FBits n
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits Natural
a)
  where n :: Int
n = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)

-- | Append 2 FBits together where the left member is shifted to make room for the right
-- element.
--
-- e.g. append (0x1 :: FBits 1) (0x70 :: FBits 7) = 0xf0 :: FBits 8
append :: forall m n r . (SizeValid m, SizeValid n, SizeValid r, (m + n) ~ r)
       => FBits n -> FBits m -> FBits r
append :: FBits n -> FBits m -> FBits r
append (FBits Natural
a) (FBits Natural
b) =
    Natural -> FBits r
forall (n :: Nat). Natural -> FBits n
FBits ((Natural
a Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`shiftL` Int
m) Natural -> Natural -> Natural
forall a. Bits a => a -> a -> a
.|.  Natural
b)
  where m :: Int
m = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy m -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy m
forall k (t :: k). Proxy t
Proxy :: Proxy m)

--appendK :: FBitsK -> FBitsK -> FBitsK
--appendK (FBitsK a) (FBitsK b) = FBitsK (a `append` b)
    -- FBits ((a `shiftL` m) .|.  b)

dropBitsOnRight :: forall a b diff . (KnownNat diff, b <= a, SizeValid a, SizeValid b, (a - b) ~ diff)
              => FBits a
              -> FBits b
dropBitsOnRight :: FBits a -> FBits b
dropBitsOnRight (FBits Natural
a) = Natural -> FBits b
forall (n :: Nat). Natural -> FBits n
FBits (Natural
a Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`shiftR` Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Proxy diff -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy diff
forall k (t :: k). Proxy t
Proxy :: Proxy diff)))

dropBitsOnLeft :: forall a b . (KnownNat b, b <= a, SizeValid a, SizeValid b)
             => FBits a
             -> FBits b
dropBitsOnLeft :: FBits a -> FBits b
dropBitsOnLeft (FBits Natural
a) = Natural -> FBits b
forall (n :: Nat). SizeValid n => Natural -> FBits n
toFBits Natural
a