path: root/protocols/codegen/Wayland/CodeGen/CodeGen.hs

{-# LANGUAGE TemplateHaskell #-}

module Wayland.CodeGen.CodeGen
  ( generateInterface
  , generateInterfaceBoot
  ) where

-- TODO: Documentation

import Wayland.CodeGen.Types
import Wayland.CodeGen.Utils
import Control.Monad.State.Strict (State, runState)
import Lens.Micro.TH (makeLenses)
import Data.List (intercalate, sort, isPrefixOf)
import Lens.Micro.Mtl ((%=))
import Data.Functor (($>))
import Control.Monad.Reader (ReaderT (..), asks)
import Data.Map ((!))
import qualified Data.Map as M
import Data.Maybe (isJust)

newtype Declaration = Declaration { declLines :: [String] }
data Import = Import { imModule :: !String
                     , imSource :: !Bool
                     }
  deriving (Show, Eq)

data GeneratorState
  = GeneratorState { _stateImports       :: ![Import]
                   , _stateExports       :: ![String]
                   , _stateExtensions    :: ![String]
                   , _stateDecls         :: ![Declaration]
                   }

data GeneratorConfig
  = GeneratorConfig { genModules   :: !FlatModules
                    , genProtocol  :: !TypedProtocol
                    , genInterface :: !InterfaceDesc
                    , genTypes     :: !TypeInfo
                    }

type Generator = ReaderT GeneratorConfig (State GeneratorState)

makeLenses ''GeneratorState

generateInterface :: FlatModules -> TypedProtocol -> TypedInterfaceDesc -> String
generateInterface flat p i = generate (GeneratorConfig flat p (fst i) (snd i)) emit

generateInterfaceBoot :: FlatModules -> TypedProtocol -> TypedInterfaceDesc -> String
generateInterfaceBoot flat p i = generate (GeneratorConfig flat p (fst i) (snd i)) emitBoot

generate :: GeneratorConfig -> Generator () -> String
generate cfg g = transform $ runState (runReaderT g cfg) initialState
  where initialState = GeneratorState { _stateImports       = []
                                      , _stateExports       = []
                                      , _stateExtensions    = []
                                      , _stateDecls         = []
                                      }
        transform (_, state) = undecls $ [ extensionDecl state, moduleDecl state, importDecl state ] ++ _stateDecls state
        extensionDecl s = Declaration $ sort $ fmap formatExtension (_stateExtensions s)
        moduleDecl s = Declaration [ "module " ++ genFullModule cfg ++ " (" ++ intercalate ", " (_stateExports s) ++ ") where" ]
        importDecl s = Declaration $ sort $ fmap formatImport (_stateImports s)
        undecls decls = intercalate "\n" $ unlines . declLines <$> filter (not . null . declLines) decls
        formatExtension e = "{-# LANGUAGE " ++ e ++ " #-}"
        formatImport i = "import " ++ (if imSource i then "{-# SOURCE #-} " else "") ++ "qualified " ++ imModule i

genFullModule :: GeneratorConfig -> String
genFullModule cfg = fullModule (genProtocol cfg) (genInterface cfg)

emitDecl :: [String] -> Generator ()
emitDecl c = stateDecls %= (++ [ Declaration c ])

emitImport :: Import -> Generator ()
emitImport i = stateImports %= \is -> if i `elem` is then is else i:is

emitExport :: String -> Generator ()
emitExport e = stateExports %= (e:)

emitExtension :: String -> Generator ()
emitExtension e = stateExtensions %= (e:)

resolve :: String -> String -> Generator String
resolve m e = do
  this <- asks genFullModule
  if   this == m
  then pure e
  else emitImport (Import m ("Wayland.Protocols." `isPrefixOf` m)) $> m ++ "." ++ e

resolveInterface :: Maybe String -> String -> Generator String
resolveInterface (Just name) e = do modules <- asks genModules
                                    let (p, i) = modules ! name
                                    resolve (fullModule p (fst i)) e
resolveInterface Nothing     e = pure e

resolveEnum :: String -> Generator String
resolveEnum e = resolveQualified $ break (== '.') e
  where resolveQualified :: (String, String) -> Generator String
        resolveQualified (e', [])   = do types <- asks genTypes
                                         pure $ typeEnums types ! e'
        resolveQualified (i,  _:e') = do modules <- asks genModules
                                         resolveInterface (Just i) $ typeEnums (snd (snd (modules ! i))) ! e'

resolveArgument :: ArgumentDesc -> Generator String
resolveArgument a = case argType a of
                      ArgumentInt    -> maybe (pure "Int")  resolveEnum (argEnum a)
                      ArgumentUInt   -> maybe (pure "Word") resolveEnum (argEnum a)
                      ArgumentFixed  -> pure "Float"
                      ArgumentArray  -> resolve "Data.ByteString" "ByteString"
                      ArgumentFD     -> resolve "System.Posix.Types" "Fd"
                      ArgumentString -> pure $ if argAllowNull a then "(Maybe String)" else "String"
                      ArgumentNewID  -> do obj <- resolveInterface (argInterface a) "Object"
                                           newObj <- resolve "Wayland.Server" "NewObject"
                                           pure $ "(" ++ newObj ++ " " ++ obj ++ ")"
                      ArgumentObject -> do obj <- resolveInterface (argInterface a) "Object"
                                           pure $ if argAllowNull a then "(Maybe " ++ obj ++ ")" else obj

emitBoot :: Generator ()
emitBoot = do
  anyObject     <- resolve "Wayland.Server" "AnyObject"
  interfaceDesc <- resolve "Wayland.Server" "InterfaceDesc"
  objectClass   <- resolve "Wayland.Server" "Object"
  enumTypes     <- asks $ typeEnums . genTypes
  i             <- asks genInterface
  emitDecl [ "newtype Object = Object " ++ anyObject ]
  emitExport "Object"
  emitDecl [ "interface :: " ++ interfaceDesc ]
  emitExport "interface"
  emitDecl [ "instance " ++ objectClass ++ " Object"
           , "instance Show Object"
           , "instance Eq Object" ]
  let emitEnumBoot e = do let name = enumTypes ! enumName e
                          emitDecl [ "data " ++ name
                                   , "instance Show " ++ name
                                   , "instance Eq " ++ name
                                   , "instance Enum " ++ name ]
                          emitExport name
  mapM_ emitEnumBoot (ifEnums i)

emit :: Generator ()
emit = do
  anyObject     <- resolve "Wayland.Server" "AnyObject"
  hashable      <- resolve "Data.Hashable" "Hashable"
  objectClass   <- resolve "Wayland.Server" "Object"
  i             <- asks genInterface
  types         <- asks genTypes
  emitExtension "GeneralizedNewtypeDeriving"
  emitExtension "EmptyDataDeriving"
  emitDecl [ "newtype Object = Object " ++ anyObject
           , "  deriving (Show, Eq, Ord, " ++ hashable ++ ")" ]
  emitExport "Object"
  mapM_ (emitMessageDesc requestDescName) (ifRequests i)
  mapM_ (emitMessageDesc eventDescName)   (ifEvents i)
  emitInterfaceDesc
  emitRequests (ifRequests i)
  emitEvents   (ifEvents   i)
  mapM_ emitEnum (ifEnums i)
  mapM_ emitArguments (typeRequestArgs types)
  mapM_ emitArguments (typeEventArgs   types)
  emitExtension "TypeFamilies"
  let matchArgument m msgTypes = case msgTypes ! msgName m of
                                   ArgumentModeNone       -> ""
                                   ArgumentModeSingle _   -> " _"
                                   ArgumentModeRecord _ _ -> " _"
  let requestSignature m =
        "  requestSignature _ (Req" ++ pascalCase (msgName m) ++ " _" ++ matchArgument m (typeRequestArgs types) ++ ") = req" ++ pascalCase (msgName m)
  let eventSignature m =
        "  eventSignature _ (Ev" ++ pascalCase (msgName m) ++ matchArgument m (typeEventArgs types) ++ ") = ev" ++ pascalCase (msgName m)
  demarshalRequests <- concat <$> mapM demarshalRequest (ifRequests i)
  marshalEvents     <- concat <$> mapM marshalEvent     (ifEvents   i)
  emitDecl ([ "instance " ++ objectClass ++ " Object where"
            , "  type ObjectEvent Object = Event"
            , "  type ObjectRequest Object = Request"
            , "  toAny (Object o) = o"
            , "  fromAny = Object"
            , "  describeInterface _ = interface" ]
            ++ demarshalRequests ++
            [ "  demarshalRequest _ _ = undefined" ]
            ++ marshalEvents ++
            [ "  marshalEvent _ _ = undefined" ]
            ++ fmap requestSignature (ifRequests i) ++
            [ "  requestSignature _ _ = undefined" ]
            ++ fmap eventSignature (ifEvents i) ++
            [ "  eventSignature _ _ = undefined" ])

requestDescName :: MessageDesc -> String
requestDescName m = "req" ++ pascalCase (msgName m)

eventDescName :: MessageDesc -> String
eventDescName m = "ev" ++ pascalCase (msgName m)

emitMessageDesc :: (MessageDesc -> String) -> MessageDesc -> Generator ()
emitMessageDesc name m = do
  let argumentNullability True  = resolve "Wayland.Server" "Nullable"
      argumentNullability False = resolve "Wayland.Server" "NotNull"
  let argumentType :: ArgumentDesc -> Generator String
      argumentType a = case argType a of
                         ArgumentInt    -> resolve "Wayland.Server" "ArgumentInt"
                         ArgumentUInt   -> resolve "Wayland.Server" "ArgumentUInt"
                         ArgumentFixed  -> resolve "Wayland.Server" "ArgumentFixed"
                         ArgumentArray  -> resolve "Wayland.Server" "ArgumentArray"
                         ArgumentFD     -> resolve "Wayland.Server" "ArgumentFD"
                         ArgumentString -> do aType        <- resolve "Wayland.Server" "ArgumentString"
                                              aNullability <- argumentNullability (argAllowNull a)
                                              pure $ aType ++ " " ++ aNullability
                         ArgumentObject -> do aType        <- resolve "Wayland.Server" "ArgumentObject"
                                              aIface       <- resolveInterface (argInterface a) "interface"
                                              aNullability <- argumentNullability (argAllowNull a)
                                              pure $ aType ++ " " ++ aNullability ++ " " ++ aIface
                         ArgumentNewID  -> do aType        <- resolve "Wayland.Server" "ArgumentNewID"
                                              aIface       <- resolveInterface (argInterface a) "interface"
                                              pure $ aType ++ " " ++ aIface
  messageDesc      <- resolve "Wayland.Server" "MessageDesc"
  msgName'         <- resolve "Wayland.Server" "msgName"
  msgArguments'    <- resolve "Wayland.Server" "msgArguments"
  msgSince'        <- resolve "Wayland.Server" "msgSince"
  msgIsDestructor' <- resolve "Wayland.Server" "msgIsDestructor"
  arguments        <- mapM argumentType (msgArguments m)
  emitDecl [ name m ++ " :: " ++ messageDesc
           , name m ++ " ="
           , "  " ++ messageDesc
           , "    { " ++ msgName' ++ " = " ++ escapeStr (msgName m)
           , "    , " ++ msgArguments' ++ " = [" ++ intercalate ", " arguments ++ "]"
           , "    , " ++ msgSince' ++ " = " ++ show (msgSince m)
           , "    , " ++ msgIsDestructor' ++ " = " ++ show (msgType m == MessageDestructor)
           , "    }"
           ]

emitInterfaceDesc :: Generator ()
emitInterfaceDesc = do
  interfaceDesc <- resolve "Wayland.Server" "InterfaceDesc"
  ifName'       <- resolve "Wayland.Server" "ifName"
  ifVersion'    <- resolve "Wayland.Server" "ifVersion"
  ifRequests'   <- resolve "Wayland.Server" "ifRequests"
  ifEvents'     <- resolve "Wayland.Server" "ifEvents"
  i             <- asks genInterface
  emitDecl [ "interface :: " ++ interfaceDesc
           , "interface ="
           , "  " ++ interfaceDesc
           , "    { " ++ ifName' ++ " = " ++ escapeStr (ifName i)
           , "    , " ++ ifVersion' ++ " = " ++ show (ifVersion i)
           , "    , " ++ ifRequests' ++ " = [" ++ intercalate ", " (requestDescName <$> ifRequests i)  ++ "]"
           , "    , " ++ ifEvents' ++ " = [" ++ intercalate ", " (eventDescName <$> ifEvents i)  ++ "]"
           , "    }" ]
  emitExport "interface"

requestConstructor :: MessageDesc -> String
requestConstructor m = "Req" ++ pascalCase (msgName m)

eventConstructor :: MessageDesc -> String
eventConstructor m = "Ev" ++ pascalCase (msgName m)

emitRequests :: [MessageDesc] -> Generator ()
emitRequests ms = emitMessages ms typeRequestArgs (\m -> requestConstructor m ++ " Object") "Request"

emitEvents :: [MessageDesc] -> Generator ()
emitEvents ms = emitMessages ms typeEventArgs eventConstructor "Event"

emitMessages :: [MessageDesc] -> (TypeInfo -> M.Map String ArgumentMode) -> (MessageDesc -> String) -> String -> Generator ()
emitMessages ms argTypes formatConstructor type' = do
  typeable     <- resolve "Data.Typeable" "Typeable"
  messageClass <- resolve "Wayland.Server" "MessageClass"
  typeClass    <- resolve "Wayland.Server" (type' ++ "Class")
  args         <- asks $ argTypes . genTypes
  let constructorArg :: ArgumentMode -> Generator String
      constructorArg ArgumentModeNone         = pure ""
      constructorArg (ArgumentModeSingle a)   = (++) " " <$> resolveArgument a
      constructorArg (ArgumentModeRecord _ r) = pure $ " " ++ r
  let constructor m = (++) (formatConstructor m) <$> constructorArg (args ! msgName m)
  constructors <- mapM constructor ms
  emitDecl $ [ "data " ++ type' ]
           ++ dataConstructors constructors
           ++ [ "  deriving (" ++ typeable ++ ", Show, Eq)" ]
  emitExport $ type' ++ "(..)"
  emitDecl [ "instance " ++ messageClass ++ " " ++ type'
           , "instance " ++ typeClass ++ " " ++ type' ]

demarshalRequest :: MessageDesc -> Generator [String]
demarshalRequest m = do
  args <- asks $ typeRequestArgs . genTypes
  let demarshalledValue :: ArgumentDesc -> Generator String
      demarshalledValue a = case argType a of
                              ArgumentInt    -> do demarshalInt' <- resolve "Wayland.Server" "demarshalInt"
                                                   if isJust (argEnum a) then pure $ "(toEnum <$> " ++ demarshalInt' ++ ")"
                                                                         else pure demarshalInt'
                              ArgumentUInt   -> do demarshalUInt' <- resolve "Wayland.Server" "demarshalUInt"
                                                   if isJust (argEnum a) then pure $ "(toEnum . fromIntegral <$> " ++ demarshalUInt' ++ ")"
                                                                         else pure demarshalUInt'
                              ArgumentFixed  -> resolve "Wayland.Server" "demarshalFixed"
                              ArgumentArray  -> resolve "Wayland.Server" "demarshalArray"
                              ArgumentFD     -> resolve "Wayland.Server" "demarshalFD"
                              ArgumentString -> if argAllowNull a then resolve "Wayland.Server" "demarshalMaybeString"
                                                                  else resolve "Wayland.Server" "demarshalString"
                              ArgumentObject -> if argAllowNull a then resolve "Wayland.Server" "demarshalMaybeObject"
                                                                  else resolve "Wayland.Server" "demarshalObject"
                              ArgumentNewID  -> resolve "Wayland.Server" "demarshalNewID"
  case args ! msgName m of
    ArgumentModeNone        -> pure [ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = pure $ " ++ requestConstructor m ++ " o" ]
    ArgumentModeSingle a    -> do val <- demarshalledValue a
                                  pure [ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = " ++ requestConstructor m ++ " o <$> " ++ val ]
    ArgumentModeRecord as r -> do vals <- mapM demarshalledValue as
                                  let demarshal (val, a) = "    " ++ camelCase (argName a) ++ "' <- " ++ val
                                  let formatArgs []      = [ "      {" ]
                                      formatArgs (a:as') =  ("      { " ++ recordField r (camelCase (argName a)) ++ " = " ++ camelCase (argName a) ++ "'") : fmap formatArg' as'
                                      formatArg' a       =   "      , " ++ recordField r (camelCase (argName a)) ++ " = " ++ camelCase (argName a) ++ "'"
                                  pure ([ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = do" ]
                                        ++ fmap demarshal (zip vals as) ++
                                        [ "    pure $ " ++ requestConstructor m ++ " o $ " ++ r ]
                                        ++ formatArgs as ++
                                        [ "      }" ])

marshalEvent :: MessageDesc -> Generator [String]
marshalEvent m = do
  args <- asks $ typeEventArgs . genTypes
  let marshalFunc :: ArgumentDesc -> Generator String
      marshalFunc a = case argType a of
                        ArgumentInt    -> do marshalInt' <- resolve "Wayland.Server" "marshalInt"
                                             if isJust (argEnum a) then pure $ "(" ++ marshalInt' ++ " . fromEnum)"
                                                                   else pure marshalInt'
                        ArgumentUInt   -> do marshalUInt' <- resolve "Wayland.Server" "marshalUInt"
                                             if isJust (argEnum a) then pure $ "(" ++ marshalUInt' ++ " . fromIntegral . fromEnum)"
                                                                   else pure marshalUInt'
                        ArgumentFixed  -> resolve "Wayland.Server" "marshalFixed"
                        ArgumentArray  -> resolve "Wayland.Server" "marshalArray"
                        ArgumentFD     -> resolve "Wayland.Server" "marshalFD"
                        ArgumentString -> if argAllowNull a then resolve "Wayland.Server" "marshalMaybeString"
                                                            else resolve "Wayland.Server" "marshalString"
                        ArgumentObject -> if argAllowNull a then resolve "Wayland.Server" "marshalMaybeObject"
                                                            else resolve "Wayland.Server" "marshalObject"
                        ArgumentNewID  -> resolve "Wayland.Server" "marshalNewID"
  case args ! msgName m of
    ArgumentModeNone        -> pure [ "  marshalEvent _ " ++ eventConstructor m ++ " = pure " ++ show (msgOpCode m) ]
    ArgumentModeSingle a    -> do func <- marshalFunc a
                                  pure [ "  marshalEvent _ (" ++ eventConstructor m ++ " x) = do"
                                       , "    " ++ func ++ " x"
                                       , "    pure " ++ show (msgOpCode m) ]
    ArgumentModeRecord as r -> do funcs <- mapM marshalFunc as
                                  let apply (func, a) = "    " ++ func ++ " $ " ++ recordField r (camelCase (argName a)) ++ " x"
                                  pure ([ "  marshalEvent _ (" ++ eventConstructor m ++ " x) = do" ]
                                        ++ fmap apply (zip funcs as) ++
                                        [ "    pure " ++ show (msgOpCode m) ])

emitEnum :: EnumDesc -> Generator ()
emitEnum e = do
  -- TODO: Handle bitsets
  types <- asks $ typeEnums . genTypes
  let name = types ! enumName e
  let constructorName entry = name ++ pascalCase (entryName entry)
  let constructorNames = fmap constructorName (enumEntries e)
  let constructors = dataConstructors constructorNames
  emitDecl ([ "data " ++ name ]
            ++ constructors ++
            [ "  deriving (Show, Eq)" ])
  emitExport $ name ++ "(..)"
  let fromEnum' (c, i) = "  fromEnum " ++ c ++ " = " ++ show i
  let toEnum' (c, i) = "  toEnum " ++ show i ++ " = " ++ c
  let constructorValues = zip constructorNames (fmap entryValue (enumEntries e))
  emitDecl ([ "instance Enum " ++ name ++ " where" ]
            ++ fmap toEnum' constructorValues ++
            [ "  toEnum _ = undefined" ]
            ++ fmap fromEnum' constructorValues)

emitArguments :: ArgumentMode -> Generator ()
emitArguments  ArgumentModeNone              = pure ()
emitArguments (ArgumentModeSingle _)         = pure ()
emitArguments (ArgumentModeRecord args name) = do
  let field a = recordField name (camelCase (argName (fst a))) ++ " :: " ++ snd a
  let formatArgs []     = [ "    {" ]
      formatArgs (a:as) =  ("    { " ++ field a) : fmap formatArg' as
      formatArg' a      =   "    , " ++ field a
  resolvedArgs <- mapM resolveArgument args
  emitDecl ([ "data " ++ name
            , "  = " ++ name ]
            ++ formatArgs (zip args resolvedArgs) ++
            [ "    }"
            , "  deriving (Show, Eq)" ])
  emitExport $ name ++ "(..)"

dataConstructors :: [String] -> [String]
dataConstructors (c:cs) = ("  = " ++ c) : fmap ("  | " ++) cs
dataConstructors []     = []

escapeStr :: String -> String
escapeStr str = "\"" ++ escapeInner str ++ "\""
  where escapeInner ('\"':xs) = "\\\"" ++ escapeInner xs
        escapeInner ('\\':xs) = "\\\\" ++ escapeInner xs
        escapeInner ('\n':xs) = "\\n"  ++ escapeInner xs
        escapeInner (x   :xs) = x      :  escapeInner xs
        escapeInner []        = []