Pure.hs

{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedLists   #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
-- Generate idiomatic haskell data types from the types in IR.
module Backends.Pure
    ( fmtModule
    ) where

import Data.Monoid                  ((<>))
import Data.String                  (IsString(..))
import GHC.Exts                     (IsList(..))
import Text.PrettyPrint.Leijen.Text (hcat, vcat)
import Text.Printf                  (printf)

import qualified Data.Map.Strict              as M
import qualified Data.Text                    as T
import qualified Text.PrettyPrint.Leijen.Text as PP

import Fmt
import IR
import Util

-- | If a module reference refers to a generated module, does it
-- refer to the raw, low-level module or the *.Pure variant (which
-- this module generates)?
data ModRefType = Pure | Raw
    deriving(Show, Read, Eq)

fmtName :: ModRefType -> Id -> Name -> PP.Doc
fmtName refTy thisMod Name{..} = modPrefix <> localName
  where
    localName = mintercalate "'" $
        map PP.textStrict $ fromList $ toList nameLocalNS ++ [nameUnqualified]
    modPrefix
        | null nsParts = ""
        | refTy == Pure && modRefToNS refTy (ByCapnpId thisMod) == ns = ""
        | otherwise = fmtModRef refTy nameModule <> "."
    ns@(Namespace nsParts) = modRefToNS refTy nameModule

modRefToNS :: ModRefType -> ModuleRef -> Namespace
modRefToNS _ (FullyQualified ns) = ns
modRefToNS ty (ByCapnpId id) = Namespace $ case ty of
    Pure -> ["Capnp", "ById", T.pack (printf "X%x" id), "Pure"]
    Raw  -> ["Capnp", "ById", T.pack (printf "X%x" id)]


fmtModule :: Module -> [(FilePath, PP.Doc)]
fmtModule mod@Module{modName=Namespace modNameParts,..} =
    [ ( T.unpack $ mintercalate "/" humanParts <> ".hs"
      , mainContent
      )
    , ( printf "Capnp/ById/X%x/Pure.hs" modId
      , vcat
            [ "{-# OPTIONS_GHC -Wno-unused-imports #-}"
            , "{-# OPTIONS_HADDOCK hide #-}"
            , "{- |"
            , hcat [ "Module: ", machineMod ]
            , hcat [ "Description: Machine-addressable alias for '", humanMod, "'." ]
            , "-}"
            , hcat [ "module ", machineMod, "(module ", humanMod, ") where" ]
            , ""
            , hcat [ "import ", humanMod ]
            ]
      )
    ]
 where
  machineMod = fmtModRef Pure (ByCapnpId modId)
  humanMod = fmtModRef Pure $ FullyQualified $ Namespace humanParts
  humanParts = "Capnp":modNameParts ++ ["Pure"]
  modFileText = PP.textStrict modFile
  mainContent = vcat
    [ "{-# LANGUAGE DuplicateRecordFields #-}"
    , "{-# LANGUAGE RecordWildCards #-}"
    , "{-# LANGUAGE FlexibleInstances #-}"
    , "{-# LANGUAGE FlexibleContexts #-}"
    , "{-# LANGUAGE MultiParamTypeClasses #-}"
    , "{-# LANGUAGE ScopedTypeVariables #-}"
    , "{-# LANGUAGE TypeFamilies #-}"
    , "{-# LANGUAGE DeriveGeneric #-}"
    , "{-# OPTIONS_GHC -Wno-unused-imports #-}"
    , "{- |"
    , "Module: " <> humanMod
    , "Description: " <> "High-level generated module for " <> modFileText
    , ""
    , "This module is the generated code for " <> modFileText <> ","
    , "for the high-level api."
    , "-}"
    , "module " <> humanMod <> " (" <> fmtExportList mod, ") where"
    , ""
    , "-- Code generated by capnpc-haskell. DO NOT EDIT."
    , "-- Generated from schema file: " <> modFileText
    , ""
    , "import Data.Int"
    , "import Data.Word"
    , ""
    , "import Data.Default (Default(def))"
    , "import GHC.Generics (Generic)"
    , ""
    , "import Data.Capnp.Basics.Pure (Data, Text)"
    , "import Control.Monad.Catch (MonadThrow)"
    , "import Data.Capnp.TraversalLimit (MonadLimit)"
    , ""
    , "import Control.Monad (forM_)"
    , ""
    , "import qualified Data.Capnp.Message as M'"
    , "import qualified Data.Capnp.Untyped as U'"
    , "import qualified Data.Capnp.Untyped.Pure as PU'"
    , "import qualified Data.Capnp.GenHelpers.Pure as PH'"
    , "import qualified Data.Capnp.Classes as C'"
    , ""
    , "import qualified Data.Vector as V"
    , "import qualified Data.ByteString as BS"
    , ""
    , fmtImport Raw $ Import (ByCapnpId modId)
    , vcat $ map (fmtImport Pure) modImports
    , vcat $ map (fmtImport Raw) modImports
    , ""
    , vcat $ map (fmtDecl modId) (M.toList modDecls)
    ]

fmtExportList :: Module -> PP.Doc
fmtExportList Module{modId,modDecls} =
    mintercalate ", " (map (fmtExport modId) (M.toList modDecls))

fmtExport :: Id -> (Name, Decl) -> PP.Doc
fmtExport thisMod (name, DeclDef DefStruct{}) =
    fmtName Pure thisMod name <> "(..)"
fmtExport thisMod (name, DeclDef DefUnion{}) =
    fmtName Pure thisMod name <> "(..)"

-- These are 'Raw' because we're just re-exporting them:
fmtExport thisMod (name, DeclDef DefEnum{}) =
    fmtName Raw thisMod name <> "(..)"
fmtExport thisMod (name, DeclConst VoidConst) =
    fmtName Raw thisMod (valueName name)
fmtExport thisMod (name, DeclConst WordConst{}) =
    fmtName Raw thisMod (valueName name)

fmtExport thisMod (name, DeclConst _) =
    fmtName Pure thisMod (valueName name)

fmtImport :: ModRefType -> Import -> PP.Doc
fmtImport ty (Import ref) = "import qualified " <> fmtModRef ty ref

fmtModRef :: ModRefType -> ModuleRef -> PP.Doc
fmtModRef ty ref = mintercalate "." (map PP.textStrict $ toList $ modRefToNS ty ref)

fmtType :: Id -> Type -> PP.Doc
fmtType thisMod (CompositeType (StructType name params)) =
    fmtName Pure thisMod name
    <> hcat [" (" <> fmtType thisMod ty <> ")" | ty <- params]
fmtType thisMod (WordType (EnumType name)) = fmtName Raw thisMod name
fmtType thisMod (PtrType (ListOf eltType)) = "PU'.ListOf (" <> fmtType thisMod eltType <> ")"
fmtType thisMod (PtrType (PtrComposite ty)) = fmtType thisMod (CompositeType ty)
fmtType _ VoidType = "()"
fmtType _ (WordType (PrimWord prim)) = fmtPrimWord prim
fmtType _ (PtrType (PrimPtr PrimText)) = "Text"
fmtType _ (PtrType (PrimPtr PrimData)) = "Data"
fmtType _ (PtrType (PrimPtr (PrimAnyPtr ty))) = "Maybe (" <> fmtAnyPtr ty <> ")"

fmtPrimWord :: PrimWord -> PP.Doc
fmtPrimWord PrimInt{isSigned=True,size}  = "Int" <> fromString (show size)
fmtPrimWord PrimInt{isSigned=False,size} = "Word" <> fromString (show size)
fmtPrimWord PrimFloat32                  = "Float"
fmtPrimWord PrimFloat64                  = "Double"
fmtPrimWord PrimBool                     = "Bool"

fmtAnyPtr :: AnyPtr -> PP.Doc
fmtAnyPtr Struct = "PU'.Struct"
fmtAnyPtr List   = "PU'.List"
fmtAnyPtr Cap    = "PU'.Cap"
fmtAnyPtr Ptr    = "PU'.PtrType"

fmtVariant :: Id -> Variant -> PP.Doc
fmtVariant thisMod Variant{variantName,variantParams} =
    fmtName Pure thisMod variantName
    <> case variantParams of
        Unnamed VoidType _ -> ""
        Unnamed ty _ -> " (" <> fmtType thisMod ty <> ")"
        Record [] -> ""
        Record fields -> PP.line <> indent
            (PP.braces $ vcat $
                PP.punctuate "," $ map (fmtField thisMod) fields)

fmtField :: Id -> Field -> PP.Doc
fmtField thisMod Field{fieldName,fieldLocType} =
    PP.textStrict fieldName <> " :: " <> fmtType thisMod fieldType
  where
    fieldType = case fieldLocType of
        VoidField      -> VoidType
        DataField _ ty -> WordType ty
        PtrField _ ty  -> PtrType ty
        HereField ty   -> CompositeType ty

fmtDecl :: Id -> (Name, Decl) -> PP.Doc
fmtDecl thisMod (name, DeclDef d)   = fmtDataDef thisMod name d
fmtDecl thisMod (name, DeclConst c) = fmtConst thisMod name c

-- | Format a constant declaration.
fmtConst :: Id -> Name -> Const -> PP.Doc
fmtConst thisMod name value =
    let pureName = fmtName Pure thisMod (valueName name)
        rawName = fmtName Raw thisMod (valueName name)
    in case value of
        -- For word types, We just re-export the definitions from the
        -- low-level module, so we don't need to declare anything here:
        VoidConst -> ""
        WordConst{} -> ""

        PtrConst{ptrType} -> vcat
            [ hcat [ pureName, " :: ", fmtType thisMod (PtrType ptrType) ]
            , hcat [ pureName, " = PH'.toPurePtrConst ", rawName ]
            ]

fmtDataDef :: Id -> Name -> DataDef -> PP.Doc
fmtDataDef thisMod dataName DefEnum{} =
    -- We end up re-exporting these, but doing nothing else.
    ""
fmtDataDef thisMod dataName dataDef =
    let rawName = fmtName Raw thisMod dataName
        pureName = fmtName Pure thisMod dataName

        unknownName = subName dataName "unknown'"
    in vcat
        [ case dataDef of
            DefEnum{} ->
                -- TODO: refactor so we don't need this case.
                error "BUG: this should have been ruled out above."
            DefStruct StructDef{fields,info} ->
                let dataVariant =
                        -- TODO: some of the functions we use still expect structs
                        -- to just be single-variant unions; this is a stand-in,
                        -- but we should update those at some point.
                        Variant
                            { variantName = dataName
                            , variantParams = Record fields
                            , variantTag = 0 -- doesn't matter; just formatting
                                             -- the data constructor, so this
                                             -- isn't used.
                            }
                in vcat
                [ data_ (fmtName Pure thisMod dataName)
                    [ fmtVariant thisMod dataVariant ]
                    ["Show", "Read", "Eq", "Generic"]
                , instance_ [] ("C'.Decerialize " <> pureName)
                    [ hcat [ "type Cerial msg ", pureName, " = ", rawName, " msg" ]
                    , "decerialize raw = do"
                    , indent $ fmtDecerializeArgs dataName fields
                    ]
                , instance_ [] ("C'.Marshal " <> pureName)
                    [ "marshalInto raw value = do"
                    , indent $ vcat
                        [ "case value of\n"
                        , indent $ vcat
                            [ fmtCerializeVariant False dataVariant ]
                        ]
                    ]
                , case info of
                    IsStandalone{} ->
                        instance_ [] ("C'.Cerialize s " <> pureName)
                            []
                    _ ->
                        ""
                ]
            DefUnion{dataVariants,parentStructName,parentStruct=StructDef{info}} ->
                let unknownVariantName = subName parentStructName "unknown'"
                in vcat
                [ data_
                    (fmtName Pure thisMod dataName)
                    (map (fmtVariant thisMod) dataVariants ++
                        [ fmtName Pure thisMod unknownVariantName <> " Word16" ]
                    )
                    ["Show", "Read", "Eq", "Generic"]
                , instance_ [] ("C'.Decerialize " <> pureName)
                    [ hcat [ "type Cerial msg ", pureName, " = ", rawName, " msg" ]
                    , "decerialize raw = do"
                    , indent $ vcat
                        [ hcat [ "raw <- ", fmtName Raw thisMod $ prefixName "get_" (subName dataName ""), " raw" ]
                        , "case raw of"
                        , indent $ vcat
                            [ vcat (map fmtDecerializeVariant dataVariants)
                            , hcat
                                [ fmtName Raw thisMod unknownName, " val -> pure $ "
                                , fmtName Pure thisMod unknownVariantName, " val"
                                ]
                            ]
                        ]
                    ]
                , instance_ [] ("C'.Marshal " <> pureName)
                    [ "marshalInto raw value = do"
                    , indent $ vcat
                        [ "case value of\n"
                        , indent $ vcat
                            [ vcat $ map (fmtCerializeVariant True) dataVariants
                            , hcat
                                [ fmtName Pure thisMod unknownVariantName, " arg_ -> "
                                , fmtName Raw thisMod $ prefixName "set_" unknownName, " raw arg_"
                                ]
                            ]
                        ]
                    ]
                , case info of
                    IsStandalone{} ->
                        instance_ [] ("C'.Cerialize s " <> pureName)
                            []
                    _ ->
                        ""
                ]
        , instance_ [] ("C'.FromStruct M'.ConstMsg " <> pureName)
            [ "fromStruct struct = do"
            , indent $ vcat
                [ "raw <- C'.fromStruct struct"
                , hcat [ "C'.decerialize (raw :: ", rawName, " M'.ConstMsg)" ]
                ]
            ]
        , instance_ [] ("Default " <> pureName)
            [ "def = PH'.defaultStruct"
            ]

        ]
  where
    fmtDecerializeArgs variantName fields = vcat
        [ hcat [ fmtName Pure thisMod variantName, " <$>" ]
        , indent $ vcat $ PP.punctuate " <*>" $
            flip map fields $ \Field{fieldName,fieldLocType} -> hcat
                [ "(", fmtName Raw thisMod $ prefixName "get_" (subName variantName fieldName)
                , " raw", case fieldLocType of
                    -- Data and void fields are always the same type in Raw and Pure forms,
                    -- so we don't need to convert them.
                    DataField _ _ -> ""
                    VoidField     -> ""
                    _             -> " >>= C'.decerialize"
                , ")"
                ]
        ]
    fmtDecerializeVariant Variant{variantName,variantParams} =
        fmtName Raw thisMod variantName <>
        case variantParams of
            Unnamed VoidType _ -> " -> pure " <> fmtName Pure thisMod variantName
            Record fields ->
              " raw -> " <> fmtDecerializeArgs variantName fields
            Unnamed (WordType _) _ -> hcat
                [ " val -> pure (", fmtName Pure thisMod variantName, " val)" ]
            _ -> hcat
                [ " val -> ", fmtName Pure thisMod variantName, " <$> C'.decerialize val" ]
    fmtCerializeVariant isUnion Variant{variantName, variantParams} =
        fmtName Pure thisMod variantName <>
        let accessorName prefix = fmtName Raw thisMod (prefixName prefix variantName)
            setterName = accessorName "set_"
        in case variantParams of
            Unnamed VoidType VoidField ->
                hcat [ " -> ", setterName, " raw" ]
            Unnamed _ (DataField _ _) ->
                hcat [ " arg_ -> ", setterName, " raw arg_"]
            Unnamed (WordType _) VoidField ->
                -- TODO: this is the unknown variant. We should find a better
                -- way to represent this; the structure of the IR here is sloppy
                -- work.
                hcat [ " arg_ -> ", setterName, " raw arg_" ]
            Unnamed _ fieldLocType -> vcat
                [ " arg_ -> do"
                , indent (fmtUseAccessors accessorName "arg_" fieldLocType)
                ]
            Record fields -> vcat
                [ "{..} -> do"
                , indent $ vcat
                    [ if isUnion
                        then "raw <- " <> setterName <> " raw"
                        else ""
                    , vcat (map (fmtCerializeField variantName) fields)
                    ]
                ]
    fmtCerializeField variantName Field{fieldName,fieldLocType} =
        let accessorName prefix = fmtName Raw thisMod $ prefixName prefix (subName variantName fieldName)
        in fmtUseAccessors accessorName fieldName fieldLocType
    fmtUseAccessors accessorName fieldVarName fieldLocType =
        let setterName = accessorName "set_"
            getterName = accessorName "get_"
            newName = accessorName "new_"
            fieldNameText = PP.textStrict fieldVarName
        in case fieldLocType of
            DataField _ _ -> hcat [ setterName, " raw ", fieldNameText ]
            VoidField -> hcat [ setterName, " raw" ]
            HereField _ -> vcat
                [ hcat [ "field_ <- ", getterName, " raw" ]
                , hcat [ "C'.marshalInto field_ ", fieldNameText ]
                ]
            PtrField _ ty -> case ty of
                PrimPtr PrimData -> vcat
                    [ hcat [ "field_ <- ", newName, " (BS.length ", fieldNameText, ") raw" ]
                    , hcat [ "C'.marshalInto field_ ", fieldNameText ]
                    ]
                PrimPtr PrimText -> vcat
                    [ hcat [ "field_ <- C'.cerialize (U'.message raw) ", fieldNameText ]
                    , hcat [ setterName, " raw field_" ]
                    ]
                PrimPtr (PrimAnyPtr _) -> vcat
                    [ hcat [ "field_ <- C'.cerialize (U'.message raw) ", fieldNameText ]
                    , hcat [ setterName, " raw field_" ]
                    ]
                ListOf eltType -> vcat
                    [ hcat [ "let len_ = V.length ", fieldNameText ]
                    , hcat [ "field_ <- ", newName, " len_ raw" ]
                    , case eltType of
                        VoidType ->
                            ""
                        CompositeType (StructType _ _) -> vcat
                            [ "forM_ [0..len_ - 1] $ \\i -> do"
                            , indent $ vcat
                                [ "elt <- C'.index i field_"
                                , hcat [ "C'.marshalInto elt (", fieldNameText, " V.! i)" ]
                                ]
                            ]
                        WordType _ -> vcat
                            [ "forM_ [0..len_ - 1] $ \\i -> do"
                            , indent $ vcat
                                [ hcat [ "C'.setIndex (", fieldNameText, " V.! i) i field_" ] ]
                            ]
                        PtrType _ ->
                            "pure ()" -- TODO
                    ]
                PtrComposite _ -> vcat
                    [ hcat [ "field_ <- ", newName, " raw" ]
                    , hcat [ "C'.marshalInto field_ ", fieldNameText ]
                    ]