GFinite.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}

-- |
-- Module      :  System.Random.GFinite
-- Copyright   :  (c) Andrew Lelechenko 2020
-- License     :  BSD-style (see the file LICENSE in the 'random' repository)
-- Maintainer  :  libraries@haskell.org
module System.Random.GFinite (
  Cardinality (..),
#if !defined(__MHS__)
  Finite (..),
  GFinite (..),
#endif /* !defined(x__MHS__) */
) where

import Data.Bits
import Data.Int
import Data.Void
import Data.Word
#if !defined(__MHS__)
import GHC.Exts (Proxy#, proxy#)
import GHC.Generics
#endif /* !defined(__MHS__) */

-- | Cardinality of a set.
data Cardinality
  = -- | Shift n is equivalent to Card (bit n)
    Shift !Int
  | Card !Integer
  deriving (Eq, Ord, Show)

-- | This is needed only as a superclass of 'Integral'.
instance Enum Cardinality where
  toEnum = fromIntegral
  fromEnum = fromIntegral
  succ = (+ 1)
  pred = subtract 1
  enumFrom x = map fromInteger (enumFrom (toInteger x))
  enumFromThen x y = map fromInteger (enumFromThen (toInteger x) (toInteger y))
  enumFromTo x y = map fromInteger (enumFromTo (toInteger x) (toInteger y))
  enumFromThenTo x y z = map fromInteger (enumFromThenTo (toInteger x) (toInteger y) (toInteger z))

instance Num Cardinality where
  fromInteger 1 = Shift 0 -- ()
  fromInteger 2 = Shift 1 -- Bool
  fromInteger n = Card n
  {-# INLINE fromInteger #-}

  x + y = fromInteger (toInteger x + toInteger y)
  {-# INLINE (+) #-}

  Shift x * Shift y = Shift (x + y)
  Shift x * Card y = Card (y `shiftL` x)
  Card x * Shift y = Card (x `shiftL` y)
  Card x * Card y = Card (x * y)
  {-# INLINE (*) #-}

  abs = Card . abs . toInteger
  signum = Card . signum . toInteger
  negate = Card . negate . toInteger

-- | This is needed only as a superclass of 'Integral'.
instance Real Cardinality where
  toRational = fromIntegral

instance Integral Cardinality where
  toInteger = \case
    Shift n -> bit n
    Card n -> n
  {-# INLINE toInteger #-}

  quotRem x' = \case
    Shift n -> (Card (x `shiftR` n), Card (x .&. (bit n - 1)))
    Card n -> let (q, r) = x `quotRem` n in (Card q, Card r)
    where
      x = toInteger x'
  {-# INLINE quotRem #-}

#if !defined(__MHS__)
-- | A type class for data with a finite number of inhabitants. This type class
-- is used in the default implementation of 'System.Random.Stateful.Uniform'.
--
-- Users are not supposed to write instances of 'Finite' manually.
-- There is a default implementation in terms of 'Generic' instead.
--
-- >>> :seti -XDeriveGeneric -XDeriveAnyClass
-- >>> import GHC.Generics (Generic)
-- >>> data MyBool = MyTrue | MyFalse deriving (Generic, Finite)
-- >>> data Action = Code MyBool | Eat (Maybe Bool) | Sleep deriving (Generic, Finite)
class Finite a where
  cardinality :: Proxy# a -> Cardinality
  toFinite :: Integer -> a
  fromFinite :: a -> Integer

  default cardinality :: (Generic a, GFinite (Rep a)) => Proxy# a -> Cardinality
  cardinality _ = gcardinality (proxy# :: Proxy# (Rep a))
  {-# INLINE cardinality #-}

  default toFinite :: (Generic a, GFinite (Rep a)) => Integer -> a
  toFinite = to . toGFinite
  {-# INLINE toFinite #-}

  default fromFinite :: (Generic a, GFinite (Rep a)) => a -> Integer
  fromFinite = fromGFinite . from
  {-# INLINE fromFinite #-}

class GFinite f where
  gcardinality :: Proxy# f -> Cardinality
  toGFinite :: Integer -> f a
  fromGFinite :: f a -> Integer

instance GFinite V1 where
  gcardinality _ = 0
  {-# INLINE gcardinality #-}
  toGFinite = const $ error "GFinite: V1 has no inhabitants"
  {-# INLINE toGFinite #-}
  fromGFinite = const $ error "GFinite: V1 has no inhabitants"
  {-# INLINE fromGFinite #-}

instance GFinite U1 where
  gcardinality _ = 1
  {-# INLINE gcardinality #-}
  toGFinite = const U1
  {-# INLINE toGFinite #-}
  fromGFinite = const 0
  {-# INLINE fromGFinite #-}

instance Finite a => GFinite (K1 _x a) where
  gcardinality _ = cardinality (proxy# :: Proxy# a)
  {-# INLINE gcardinality #-}
  toGFinite = K1 . toFinite
  {-# INLINE toGFinite #-}
  fromGFinite = fromFinite . unK1
  {-# INLINE fromGFinite #-}

instance GFinite a => GFinite (M1 _x _y a) where
  gcardinality _ = gcardinality (proxy# :: Proxy# a)
  {-# INLINE gcardinality #-}
  toGFinite = M1 . toGFinite
  {-# INLINE toGFinite #-}
  fromGFinite = fromGFinite . unM1
  {-# INLINE fromGFinite #-}

instance (GFinite a, GFinite b) => GFinite (a :+: b) where
  gcardinality _ =
    gcardinality (proxy# :: Proxy# a) + gcardinality (proxy# :: Proxy# b)
  {-# INLINE gcardinality #-}

  toGFinite n
    | n < cardA = L1 $ toGFinite n
    | otherwise = R1 $ toGFinite (n - cardA)
    where
      cardA = toInteger (gcardinality (proxy# :: Proxy# a))
  {-# INLINE toGFinite #-}

  fromGFinite = \case
    L1 x -> fromGFinite x
    R1 x -> fromGFinite x + toInteger (gcardinality (proxy# :: Proxy# a))
  {-# INLINE fromGFinite #-}

instance (GFinite a, GFinite b) => GFinite (a :*: b) where
  gcardinality _ =
    gcardinality (proxy# :: Proxy# a) * gcardinality (proxy# :: Proxy# b)
  {-# INLINE gcardinality #-}

  toGFinite n = toGFinite (toInteger q) :*: toGFinite (toInteger r)
    where
      cardB = gcardinality (proxy# :: Proxy# b)
      (q, r) = Card n `quotRem` cardB
  {-# INLINE toGFinite #-}

  fromGFinite (q :*: r) =
    toInteger (gcardinality (proxy# :: Proxy# b) * Card (fromGFinite q)) + fromGFinite r
  {-# INLINE fromGFinite #-}

instance Finite Void

instance Finite ()

instance Finite Bool

instance Finite Ordering

instance Finite Char where
  cardinality _ = Card $ toInteger (fromEnum (maxBound :: Char)) + 1
  {-# INLINE cardinality #-}
  toFinite = toEnum . fromInteger
  {-# INLINE toFinite #-}
  fromFinite = toInteger . fromEnum
  {-# INLINE fromFinite #-}

cardinalityDef :: forall a. (Num a, FiniteBits a) => Proxy# a -> Cardinality
cardinalityDef _ = Shift (finiteBitSize (0 :: a))

toFiniteDef :: forall a. (Num a, FiniteBits a) => Integer -> a
toFiniteDef n
  | isSigned (0 :: a) = fromInteger (n - bit (finiteBitSize (0 :: a) - 1))
  | otherwise = fromInteger n

fromFiniteDef :: (Integral a, FiniteBits a) => a -> Integer
fromFiniteDef x
  | isSigned x = toInteger x + bit (finiteBitSize x - 1)
  | otherwise = toInteger x

instance Finite Word8 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Word16 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Word32 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Word64 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Word where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Int8 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Int16 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Int32 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Int64 where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite Int where
  cardinality = cardinalityDef
  {-# INLINE cardinality #-}
  toFinite = toFiniteDef
  {-# INLINE toFinite #-}
  fromFinite = fromFiniteDef
  {-# INLINE fromFinite #-}

instance Finite a => Finite (Maybe a)

instance (Finite a, Finite b) => Finite (Either a b)

instance (Finite a, Finite b) => Finite (a, b)

instance (Finite a, Finite b, Finite c) => Finite (a, b, c)

instance (Finite a, Finite b, Finite c, Finite d) => Finite (a, b, c, d)

instance (Finite a, Finite b, Finite c, Finite d, Finite e) => Finite (a, b, c, d, e)

instance (Finite a, Finite b, Finite c, Finite d, Finite e, Finite f) => Finite (a, b, c, d, e, f)

instance
  (Finite a, Finite b, Finite c, Finite d, Finite e, Finite f, Finite g) =>
  Finite (a, b, c, d, e, f, g)

#endif /* !defined(__MHS__) */