Skip to content

Instantly share code, notes, and snippets.

@AndreasPK

AndreasPK/Packing

Created Jul 22, 2020
Embed
What would you like to do?
{- | Pack types of a known size into as few bytes as possible.
We do so by assigning each instance of the Packable class
a 'Width' in bits.
When combining types we add up their widths using type level computation.
-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ConstraintKinds #-}
-- Allow constraint on result of toBits.
{-# LANGUAGE ConstrainedClassMethods #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# OPTIONS_GHC -ddump-simpl -ddump-cmm -ddump-to-file -fforce-recomp #-}
-- {-# OPTIONS_GHC -dsuppress-all #-}
module Data.Packed.Static
( KnownPackable(..)
, SizedBoundedEnum(..)
, SizedInt(..)
, RepType
, BitsFitIn
, bitWidth
, setBits
)
where
-- import Internal.Util hiding (setBits)
import Data.Bits
import Data.Word
import Data.Int
import Data.Proxy
import Data.Kind
import GHC.TypeLits as T
import GHC.Exts (Proxy#, proxy#)
import Debug.Trace
-- | Allow packing of this type into a number of bits known at compile time.
class KnownPackable (rep :: Type) (a :: Type) | a -> rep where
type BitWidth a :: n
-- | Store as bits
toBits :: a -> rep
-- | Read from bits, we do not constrain this so we can eg read a Bool out of a Int or Word8
fromBits :: (BitRep bits) => bits -> a
-- | Type level `if`
type family If c t e where
If 'True t e = t
If 'False t e = e
-- | Will always return a Integral type
type family BitsFitIn (a :: Nat) :: Type where
BitsFitIn 0 = Word8
BitsFitIn (n ) =
If (n T.<=? 8) (Word8)
(If (n T.<=? 16) (Word16)
(If (n T.<=? 32) (Word32)
(If (n T.<=? 64) (Word64)
(TypeError ('Text "Bit fields >= 64 bit not supported"))
)
)
)
type family RepType (a :: Type) :: Type where
RepType a = BitsFitIn (BitWidth a)
-- type family BitWidth a :: Nat
type family Max (n1 :: Nat) (n2 :: Nat) :: Nat where
Max n1 n2 = If (n1 T.<=? n2) n2 n1
-- | Constraints required of a bitmap
type BitRep a = (Integral a, Bits a, Num a)
-- | Set the first x bits which the given instance occupies
--
{-# INLINE setBits #-}
setBits :: forall bitwidth (rep :: Type) (p :: Type). ( bitwidth ~ ((BitWidth p))
, KnownNat bitwidth
, BitRep rep
, rep ~ BitsFitIn bitwidth
)
=> Proxy# p -> rep
-- setBits _ = fromIntegral $ ((2 :: Int) ^ bitcount) - 1
setBits _ = fromIntegral $ ((1 :: Integer) `shiftL` bitcount) - 1
where
bitcount = fromIntegral $
natVal' ( proxy# :: Proxy# bitwidth ) :: Int
-- | Get the number of bits p occupies when packed.
{-# INLINE bitWidth #-}
bitWidth :: forall a width rep. (width ~ BitWidth a, rep ~ RepType a, BitRep rep, KnownNat width)
=> Proxy# a -> Int
bitWidth _ = fromIntegral $ natVal' ( proxy# :: Proxy# (BitWidth a) ) :: Int
instance (rep ~ RepType Bool) => KnownPackable rep Bool where
type BitWidth Bool = 1
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits x = fromIntegral . fromEnum $ x
fromBits bits = toEnum tag
where
tag = fromIntegral masked :: Int
masked = fromIntegral bits .&. range
range = setBits (proxy# :: Proxy# Bool)
instance forall rep repa a. ( RepType (Maybe a) ~ rep
, KnownPackable repa a
, BitRep rep
, BitRep repa
) => KnownPackable rep (Maybe a) where
type BitWidth (Maybe a) = BitWidth a + 1
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits Nothing = fromIntegral (0 :: Int) :: rep
toBits (Just x) = 1 .|. (fromIntegral (toBits x) `unsafeShiftL` 1)
fromBits bits
| not (testBit bits 0) = Nothing
| otherwise = Just . fromBits $ bits `unsafeShiftR` 1
instance forall rep repa repb a b.
( rep ~ BitsFitIn (BitWidth (a,b)) , repa ~ RepType a, repb ~ RepType b
, KnownNat (BitWidth a), KnownNat (BitWidth b)
, KnownPackable repa a, KnownPackable repb b
, BitRep rep, BitRep repa, BitRep repb
, Show rep
) => KnownPackable rep (a,b) where
type BitWidth (a,b) = BitWidth a + BitWidth b
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits (a,b) = let widthA = fromIntegral $ bitWidth (proxy# :: Proxy# a)
bitsA = fromIntegral (toBits a) :: rep
bitsB = fromIntegral (toBits b) :: rep
in bitsA .|. (bitsB `unsafeShiftL` widthA)
fromBits bits =
let widthA = fromIntegral $ bitWidth (proxy# :: Proxy# a)
a = fromBits bits
b = fromBits (bits `unsafeShiftR` widthA)
in (a,b)
instance forall rep repa repb repc a b c.
( rep ~ RepType (a,b,c)
, repa ~ RepType a, repb ~ RepType b, repc ~ RepType c
, KnownNat (BitWidth a), KnownNat (BitWidth b), KnownNat (BitWidth c)
, KnownPackable repa a, KnownPackable repb b, KnownPackable repc c
, BitRep rep, BitRep repa, BitRep repb, BitRep repc
) => KnownPackable rep (a,b,c) where
type BitWidth (a,b,c) = BitWidth a + BitWidth b + BitWidth c
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits (a,b,c) = fromIntegral (toBits a) .|.
((fromIntegral (toBits b)) `unsafeShiftL` widthA) .|.
((fromIntegral (toBits c)) `unsafeShiftL` (widthA + widthB))
where
widthA = fromIntegral $ bitWidth (proxy# :: Proxy# a)
widthB = fromIntegral $ bitWidth (proxy# :: Proxy# b)
fromBits bits =
let widthA = fromIntegral $ bitWidth (proxy# :: Proxy# a)
widthB = fromIntegral $ bitWidth (proxy# :: Proxy# b)
a = fromBits bits
b = fromBits (bits `unsafeShiftR` widthA)
c = fromBits (bits `unsafeShiftR` (widthA + widthB))
in (a,b,c)
-----------------------------------------
-- Useful newtype wrappers
-----------------------------------------
-- | Derive a Packed instance based on 'Bounded','Enum' and the bits required for storage.
newtype SizedBoundedEnum (size :: Nat) a = SizedBoundedEnum a deriving (Show,Bounded,Enum)
instance forall a rep size. (Bounded a, Enum a, rep ~ BitsFitIn size, BitRep rep
,KnownNat size
,KnownNat (size + 1)
,BitRep rep)
=> KnownPackable rep (SizedBoundedEnum size a) where
type BitWidth (SizedBoundedEnum size a) = size
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits (SizedBoundedEnum x) = fromIntegral . fromEnum $ x
fromBits bits = SizedBoundedEnum $ toEnum tag
where
tag = fromIntegral masked :: Int
masked = fromIntegral bits .&. range
range = setBits (proxy# :: Proxy# (SizedBoundedEnum size a))
-- | Only store 'bitcount' bits.
newtype SizedInt (bitcount :: Nat) = SizedInt { getInt :: Int } deriving (Eq,Ord,Show,Num,Enum,Integral,Real)
instance ( KnownNat size, (BitsFitIn size) ~ rep
, BitRep rep)
=> KnownPackable rep (SizedInt size) where
type BitWidth (SizedInt size) = size
{-# INLINE toBits #-}
{-# INLINE fromBits #-}
toBits (SizedInt x) = fromIntegral x
fromBits bits = SizedInt value
where
value = fromIntegral masked :: Int
masked = fromIntegral bits .&. range
range = setBits (proxy# :: Proxy# (SizedInt size))
deriving via SizedBoundedEnum 8 Word8 instance KnownPackable Word8 Word8
deriving via SizedBoundedEnum 16 Word16 instance KnownPackable Word16 Word16
deriving via SizedBoundedEnum 32 Word32 instance KnownPackable Word32 Word32
deriving via SizedBoundedEnum 64 Word64 instance KnownPackable Word64 Word64
-- deriving via SizedBoundedEnum 8 Int8 instance KnownPackable Word8 Int8
deriving via SizedBoundedEnum 16 Int16 instance KnownPackable Word16 Int16
deriving via SizedBoundedEnum 32 Int32 instance KnownPackable Word32 Int32
deriving via SizedBoundedEnum 64 Int64 instance KnownPackable Word64 Int64
-- data Foo = X | Y deriving (Bounded,Enum,Show)
-- deriving via SizedBoundedEnum 1 Foo instance KnownPackable Word8 Foo
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment