EnumGenerator.hs

-- |
-- Module:      Parser.EnumGenerator
-- Copyright:   (c) 2015-2016 Martijn Rijkeboer <mrr@sru-systems.com>
-- License:     MIT
-- Maintainer:  Martijn Rijkeboer <mrr@sru-systems.com>
--
-- Generate code for Enum types.

module Parser.EnumGenerator
    ( getEnumCode
    ) where

import Control.Monad.State (State)
import Data.List (foldl', groupBy, intersperse, sortOn)
import Data.Monoid ((<>))
import Data.Text.Lazy (Text)
import Data.Text.Lazy.Builder (Builder, fromString, toLazyText)
import Parser.EnumDesc (EnumDesc)
import Parser.EnumValueDesc (EnumValueDesc)
import Parser.FileDesc (FileDesc)
import Parser.GeneratorUtils

import qualified Parser.EnumDesc       as EnumDesc
import qualified Parser.EnumValueDesc  as EnumValueDesc
import qualified Parser.FileDesc       as FileDesc


getEnumCode :: FileDesc -> EnumDesc -> State GenState Text
getEnumCode f ed = do
    commentLine <- getComment
    moduleLine <- getModule f ed
    importLines <- getImports f ed
    typeLines <- getType f ed
    defaultInst <- getDefaultInst ed
    mergeableInst <- getMergeableInst ed
    wireEnumInst <- getWireEnumInst ed
    return $ toLazyText $
      commentLine <> nl <>
      moduleLine <> nl <>
      importLines <> nl <>
      typeLines <> nl <>
      defaultInst <> nl <>
      mergeableInst <> nl <>
      wireEnumInst <> nl


getModule :: FileDesc -> EnumDesc -> State GenState Builder
getModule f ed = do
    modName <- getModuleName f ed
    return $
      fromString "module " <>
      fromString modName <>
      fromString " where" <> nl


getModuleName :: FileDesc -> EnumDesc -> State GenState String
getModuleName f ed = return $ case FileDesc.getPackage f of
    Just package -> package ++ "." ++ enumName
    Nothing      -> fileName ++ "." ++ enumName
  where
    fileName = FileDesc.getName f
    enumName = EnumDesc.getName ed


getImports :: FileDesc -> EnumDesc -> State GenState Builder
getImports _ _ = return $
    getUnqualifiedImport "Prelude" "" <>
    getQualifiedAsImport "Data.ProtoBufInt" "PB"


getType :: FileDesc -> EnumDesc -> State GenState Builder
getType f ed = do
    enumValues <- getEnumValues f ed
    return $
      fromString "data " <>
      fromString name <>
      fromString " = " <>
      enumValues <>
      nl <>
      tab <>
      fromString "deriving (PB.Show, PB.Eq, PB.Ord)" <> nl
  where name = EnumDesc.getName ed


getEnumValues :: FileDesc -> EnumDesc -> State GenState Builder
getEnumValues _ ed = fmap (foldl' (<>) empty) ivalues
  where
    ivalues = fmap (intersperse separator) fvalues
    fvalues = mapM getEnumValue (EnumDesc.getValueDescs ed)
    separator = fromString " | "


getEnumValue :: EnumValueDesc -> State GenState Builder
getEnumValue = return . fromString . EnumValueDesc.getName


getDefaultInst :: EnumDesc -> State GenState Builder
getDefaultInst ed = return $
    fromString "instance PB.Default " <> fromString name <> fromString " where" <> nl <>
    tab <> fromString "defaultVal = " <> fromString first <> nl
  where
    name = EnumDesc.getName ed
    first = EnumValueDesc.getName $ head $ EnumDesc.getValueDescs ed


getMergeableInst :: EnumDesc -> State GenState Builder
getMergeableInst ed = return $
    fromString "instance PB.Mergeable " <> fromString name <> fromString " where" <> nl
  where
    name = EnumDesc.getName ed


getWireEnumInst :: EnumDesc -> State GenState Builder
getWireEnumInst ed = return $
    fromString "instance PB.WireEnum " <> fromString name <> fromString " where" <> nl <>
    intToEnumLines <> nl <>
    enumToIntLines <> nl
  where
    name = EnumDesc.getName ed
    intToEnumLines = getIntToEnumLines (EnumDesc.getValueDescs ed)
    enumToIntLines = getEnumToIntLines (EnumDesc.getValueDescs ed)


getIntToEnumLines :: [EnumValueDesc] -> Builder
getIntToEnumLines es =
    foldl' (<>) empty (map getIntToEnumLine $ uniqEs es) <>
    tab <> fromString "intToEnum _ = PB.defaultVal" <> nl
  where
    uniqEs = map head . groupBy areEq . sortOn EnumValueDesc.getNumber
    areEq a b = EnumValueDesc.getNumber a == EnumValueDesc.getNumber b


getIntToEnumLine :: EnumValueDesc -> Builder
getIntToEnumLine e = tab <> fromString "intToEnum " <> number <> fromString " = " <> name <> nl
  where
    name = fromString $ EnumValueDesc.getName e
    number = fromString $ show $ EnumValueDesc.getNumber e


getEnumToIntLines :: [EnumValueDesc] -> Builder
getEnumToIntLines es = foldl' (<>) empty (map getEnumToIntLine es)


getEnumToIntLine :: EnumValueDesc -> Builder
getEnumToIntLine e = tab <> fromString "enumToInt " <> name <> fromString " = " <> number <> nl
  where
    name = fromString $ EnumValueDesc.getName e
    number = fromString $ show $ EnumValueDesc.getNumber e