ByteEnum.hs

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE TemplateHaskell       #-}

module ByteEnum (ByteEnum, toByte, fromByte, fromByteAll, allValues, enum) where

import           Data.Word           (Word8)
import           Language.Haskell.TH

class ByteEnum e where
  toByte :: e -> Word8
  fromByte :: Word8 -> Maybe e
  allValues :: [e]

fromByteAll :: ByteEnum e => Word8 -> [e]
fromByteAll b = filter (\v -> toByte v == b) allValues

instance ByteEnum e => Enum (Maybe e) where
  toEnum x
    | x >= 0 && x < 256 = fromByte $ fromIntegral x
    | otherwise = Nothing
  fromEnum (Just e) = fromIntegral $ toByte e
  fromEnum Nothing  = 0

fromByteQ :: Name
fromByteQ = mkName "fromByte"

toByteQ :: Name
toByteQ = mkName "toByte"

allValuesQ :: Name
allValuesQ = mkName "allValues"

-- | Generate enum with byte value associated to each enum value.
--
-- Example:
-- >>> enum "SomeEnum" ["Show", "Eq"] [("Foo", 1), ("Bar", 2), ("Baz", 3)]
-- data SomeEnum = Foo | Bar | Baz deriving (Show, Eq)
-- instance ByteEnum SomeEnum where
--   toByte Foo = 1
--   toByte Bar = 2
--   toByte Baz = 3
--   fromByte 1 = Just Foo
--   fromByte 2 = Just Bar
--   fromByte 3 = Just Baz
--   fromByte _ = Nothing
enum :: String -> [String] -> [(String, Word8)] -> Q [Dec]
enum name ds cs = do
  instD <- instanceD (cxt []) (appT [t|ByteEnum|] (conT name'))
    [ fromByteD
    , toByteD
    , allValuesD
    ]
  return [enumD, instD]
    where
      name' = mkName name
      csNames = map (mkName . fst) cs
      csNames' = map (`NormalC` []) csNames
      enumD = DataD [] name' [] Nothing csNames' [DerivClause Nothing $ map (ConT . mkName) ds]
      fromByteD = funD toByteQ $ map toByteClause cs
      toByteD = funD fromByteQ $ map fromByteClause cs ++ [defaultFromEnumClause]
      allValuesD = funD allValuesQ [clause [] (normalB $ listE $ map conE csNames) []]

      defaultFromEnumClause :: ClauseQ
      defaultFromEnumClause = clause [wildP] (normalB [|Nothing|]) []

      toByteClause :: (String, Word8) -> ClauseQ
      toByteClause (c, n) = clause
        [conP (mkName c) []]
        (normalB $ litE $ integerL $ fromIntegral n)
        []

      fromByteClause :: (String, Word8) -> ClauseQ
      fromByteClause (c, n) = clause
        [litP $ integerL $ fromIntegral n]
        (normalB [| Just $(conE $ mkName c) |])
        []