{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Data.Function.FastMemo.Word () where

import Data.Bits
import Data.Foldable (foldl')
import Data.Function.FastMemo.Class (Memoizable (..))
import Data.Function.FastMemo.Util (memoizeFixedLen)
import qualified Data.Vector as V
import Data.Word

instance Memoizable Word8 where
  memoize :: (Word8 -> b) -> Word8 -> b
memoize Word8 -> b
f =
    let values :: Vector b
values = Word8 -> b
f (Word8 -> b) -> Vector Word8 -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Word8] -> Vector Word8
forall a. [a] -> Vector a
V.fromList [Word8
0x00 .. Word8
0xff]
     in \Word8
i -> Vector b
values Vector b -> Int -> b
forall a. Vector a -> Int -> a
V.! Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
i

deriving via MemoWord Word instance Memoizable Word

deriving via MemoWord Word16 instance Memoizable Word16

deriving via MemoWord Word32 instance Memoizable Word32

deriving via MemoWord Word64 instance Memoizable Word64

newtype MemoWord a = MemoWord {MemoWord a -> a
getMemoWord :: a}

instance (FiniteBits a, Integral a) => Memoizable (MemoWord a) where
  memoize :: (MemoWord a -> b) -> MemoWord a -> b
memoize MemoWord a -> b
f =
    Int -> ([Word8] -> b) -> [Word8] -> b
forall a b.
(HasCallStack, Memoizable a) =>
Int -> ([a] -> b) -> [a] -> b
memoizeFixedLen (a -> Int
forall a. FiniteBits a => a -> Int
byteLen (a
0 :: a)) (MemoWord a -> b
f (MemoWord a -> b) -> ([Word8] -> MemoWord a) -> [Word8] -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> MemoWord a
forall a. a -> MemoWord a
MemoWord (a -> MemoWord a) -> ([Word8] -> a) -> [Word8] -> MemoWord a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> a
forall a. (Bits a, Num a) => [Word8] -> a
fromBytes) ([Word8] -> b) -> (MemoWord a -> [Word8]) -> MemoWord a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [Word8]
forall a. (FiniteBits a, Integral a) => a -> [Word8]
toBytes (a -> [Word8]) -> (MemoWord a -> a) -> MemoWord a -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoWord a -> a
forall a. MemoWord a -> a
getMemoWord

byteLen :: FiniteBits a => a -> Int
byteLen :: a -> Int
byteLen a
x = (a -> Int
forall a. FiniteBits a => a -> Int
finiteBitSize a
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
8

toBytes :: (FiniteBits a, Integral a) => a -> [Word8]
toBytes :: a -> [Word8]
toBytes a
x = [a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
i) | Int
i <- [Int
s0, Int
s0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
8 .. Int
0]]
  where
    s0 :: Int
s0 = (a -> Int
forall a. FiniteBits a => a -> Int
byteLen a
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8

fromBytes :: (Bits a, Num a) => [Word8] -> a
fromBytes :: [Word8] -> a
fromBytes = (a -> Word8 -> a) -> a -> [Word8] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\a
acc Word8
x -> a
acc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
8 a -> a -> a
forall a. Bits a => a -> a -> a
.|. Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
x) a
0