{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE NoFieldSelectors #-}

module Elara.Core.Module where

import Data.Generics.Product
import Elara.AST.Name (ModuleName, Qualified)
import Elara.Core (CoreBind, DataCon, Type, TypeVariable)
import Elara.Core qualified as Core
import Elara.Core.ANF qualified as ANF
import Elara.Core.Pretty (prettyTy, prettyTypeVariables)
import Elara.Data.Kind (ElaraKind)
import Elara.Data.Pretty (AnsiStyle, Doc, Pretty (pretty), bracedBlock, hardline, indentDepth, nest, (<+>))
import Elara.Data.Pretty.Styles (keyword)
import Elara.Data.TopologicalGraph (HasDependencies (..))
import Elara.Pretty.Common (prettyCtorsInline)

data CoreModule bind = CoreModule
    { forall bind. CoreModule bind -> ModuleName
name :: !ModuleName
    , forall bind. CoreModule bind -> [CoreDeclaration bind]
declarations :: ![CoreDeclaration bind]
    }
    deriving ((forall x. CoreModule bind -> Rep (CoreModule bind) x)
-> (forall x. Rep (CoreModule bind) x -> CoreModule bind)
-> Generic (CoreModule bind)
forall x. Rep (CoreModule bind) x -> CoreModule bind
forall x. CoreModule bind -> Rep (CoreModule bind) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall bind x. Rep (CoreModule bind) x -> CoreModule bind
forall bind x. CoreModule bind -> Rep (CoreModule bind) x
$cfrom :: forall bind x. CoreModule bind -> Rep (CoreModule bind) x
from :: forall x. CoreModule bind -> Rep (CoreModule bind) x
$cto :: forall bind x. Rep (CoreModule bind) x -> CoreModule bind
to :: forall x. Rep (CoreModule bind) x -> CoreModule bind
Generic)

-- the constraint is necessary for 'gplate' to be able to find the fields correctly
instance HasDependencies (CoreModule CoreBind) where
    type Key (CoreModule CoreBind) = ModuleName
    key :: CoreModule CoreBind -> Key (CoreModule CoreBind)
key = Optic' A_Lens NoIx (CoreModule CoreBind) ModuleName
-> CoreModule CoreBind -> ModuleName
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view (forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"name")

    dependencies :: CoreModule CoreBind -> [Key (CoreModule CoreBind)]
dependencies CoreModule CoreBind
m = do
        CoreModule CoreBind
m CoreModule CoreBind
-> Optic' A_Traversal NoIx (CoreModule CoreBind) ModuleName
-> [ModuleName]
forall k s (is :: IxList) a.
Is k A_Fold =>
s -> Optic' k is s a -> [a]
^.. forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"declarations" Lens
  (CoreModule CoreBind)
  (CoreModule CoreBind)
  [CoreDeclaration CoreBind]
  [CoreDeclaration CoreBind]
-> Optic
     A_Traversal
     NoIx
     [CoreDeclaration CoreBind]
     [CoreDeclaration CoreBind]
     (Qualified Text)
     (Qualified Text)
-> Optic
     A_Traversal
     NoIx
     (CoreModule CoreBind)
     (CoreModule CoreBind)
     (Qualified Text)
     (Qualified Text)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% (forall a s. (GPlate a s, SafeGPlate (Rep s) a) => Traversal' s a
genericPlate @(Qualified Text)) Optic
  A_Traversal
  NoIx
  (CoreModule CoreBind)
  (CoreModule CoreBind)
  (Qualified Text)
  (Qualified Text)
-> Optic
     A_Lens NoIx (Qualified Text) (Qualified Text) ModuleName ModuleName
-> Optic' A_Traversal NoIx (CoreModule CoreBind) ModuleName
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"qualifier"

instance HasDependencies (CoreModule (ANF.TopLevelBind Core.Var)) where
    type Key (CoreModule (ANF.TopLevelBind Core.Var)) = ModuleName
    key :: CoreModule (TopLevelBind Var)
-> Key (CoreModule (TopLevelBind Var))
key = Optic' A_Lens NoIx (CoreModule (TopLevelBind Var)) ModuleName
-> CoreModule (TopLevelBind Var) -> ModuleName
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view (forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"name")

    dependencies :: CoreModule (TopLevelBind Var)
-> [Key (CoreModule (TopLevelBind Var))]
dependencies CoreModule (TopLevelBind Var)
m = do
        CoreModule (TopLevelBind Var)
m CoreModule (TopLevelBind Var)
-> Optic'
     A_Traversal NoIx (CoreModule (TopLevelBind Var)) ModuleName
-> [ModuleName]
forall k s (is :: IxList) a.
Is k A_Fold =>
s -> Optic' k is s a -> [a]
^.. forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"declarations" Lens
  (CoreModule (TopLevelBind Var))
  (CoreModule (TopLevelBind Var))
  [CoreDeclaration (TopLevelBind Var)]
  [CoreDeclaration (TopLevelBind Var)]
-> Optic
     A_Traversal
     NoIx
     [CoreDeclaration (TopLevelBind Var)]
     [CoreDeclaration (TopLevelBind Var)]
     (Qualified Text)
     (Qualified Text)
-> Optic
     A_Traversal
     NoIx
     (CoreModule (TopLevelBind Var))
     (CoreModule (TopLevelBind Var))
     (Qualified Text)
     (Qualified Text)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% (forall a s. (GPlate a s, SafeGPlate (Rep s) a) => Traversal' s a
genericPlate @(Qualified Text)) Optic
  A_Traversal
  NoIx
  (CoreModule (TopLevelBind Var))
  (CoreModule (TopLevelBind Var))
  (Qualified Text)
  (Qualified Text)
-> Optic
     A_Lens NoIx (Qualified Text) (Qualified Text) ModuleName ModuleName
-> Optic'
     A_Traversal NoIx (CoreModule (TopLevelBind Var)) ModuleName
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% forall (field :: Symbol) s t a b.
HasField field s t a b =>
Lens s t a b
field @"qualifier"

data CoreDeclaration bind
    = CoreValue bind
    | CoreType CoreTypeDecl
    deriving ((forall x. CoreDeclaration bind -> Rep (CoreDeclaration bind) x)
-> (forall x. Rep (CoreDeclaration bind) x -> CoreDeclaration bind)
-> Generic (CoreDeclaration bind)
forall x. Rep (CoreDeclaration bind) x -> CoreDeclaration bind
forall x. CoreDeclaration bind -> Rep (CoreDeclaration bind) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall bind x. Rep (CoreDeclaration bind) x -> CoreDeclaration bind
forall bind x. CoreDeclaration bind -> Rep (CoreDeclaration bind) x
$cfrom :: forall bind x. CoreDeclaration bind -> Rep (CoreDeclaration bind) x
from :: forall x. CoreDeclaration bind -> Rep (CoreDeclaration bind) x
$cto :: forall bind x. Rep (CoreDeclaration bind) x -> CoreDeclaration bind
to :: forall x. Rep (CoreDeclaration bind) x -> CoreDeclaration bind
Generic)

data CoreTypeDecl = CoreTypeDecl
    { CoreTypeDecl -> Qualified Text
ctdName :: !(Qualified Text)
    , CoreTypeDecl -> ElaraKind
kind :: !ElaraKind
    , CoreTypeDecl -> [TypeVariable]
typeVars :: ![TypeVariable]
    , CoreTypeDecl -> CoreTypeDeclBody
typeBody :: CoreTypeDeclBody
    }
    deriving ((forall x. CoreTypeDecl -> Rep CoreTypeDecl x)
-> (forall x. Rep CoreTypeDecl x -> CoreTypeDecl)
-> Generic CoreTypeDecl
forall x. Rep CoreTypeDecl x -> CoreTypeDecl
forall x. CoreTypeDecl -> Rep CoreTypeDecl x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CoreTypeDecl -> Rep CoreTypeDecl x
from :: forall x. CoreTypeDecl -> Rep CoreTypeDecl x
$cto :: forall x. Rep CoreTypeDecl x -> CoreTypeDecl
to :: forall x. Rep CoreTypeDecl x -> CoreTypeDecl
Generic)

data CoreTypeDeclBody
    = CoreTypeAlias Type
    | CoreDataDecl Core.TyCon [DataCon]
    deriving ((forall x. CoreTypeDeclBody -> Rep CoreTypeDeclBody x)
-> (forall x. Rep CoreTypeDeclBody x -> CoreTypeDeclBody)
-> Generic CoreTypeDeclBody
forall x. Rep CoreTypeDeclBody x -> CoreTypeDeclBody
forall x. CoreTypeDeclBody -> Rep CoreTypeDeclBody x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CoreTypeDeclBody -> Rep CoreTypeDeclBody x
from :: forall x. CoreTypeDeclBody -> Rep CoreTypeDeclBody x
$cto :: forall x. Rep CoreTypeDeclBody x -> CoreTypeDeclBody
to :: forall x. Rep CoreTypeDeclBody x -> CoreTypeDeclBody
Generic)

instance Pretty bind => Pretty (CoreModule bind) where
    pretty :: CoreModule bind -> Doc AnsiStyle
pretty (CoreModule ModuleName
name [CoreDeclaration bind]
decls) =
        Doc AnsiStyle
"module"
            Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ModuleName -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty ModuleName
name
            Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Doc AnsiStyle
forall ann. Doc ann
hardline
            Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Int -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Int -> Doc ann -> Doc ann
nest Int
indentDepth ([CoreDeclaration bind] -> Doc AnsiStyle
forall a. Pretty a => [a] -> Doc AnsiStyle
bracedBlock [CoreDeclaration bind]
decls)

instance Pretty bind => Pretty (CoreDeclaration bind) where
    pretty :: CoreDeclaration bind -> Doc AnsiStyle
pretty (CoreValue bind
v) = bind -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty bind
v
    pretty (CoreType CoreTypeDecl
t) = CoreTypeDecl -> Doc AnsiStyle
prettyTdef CoreTypeDecl
t

instance Pretty CoreTypeDecl where
    pretty :: CoreTypeDecl -> Doc AnsiStyle
pretty = CoreTypeDecl -> Doc AnsiStyle
prettyTdef

instance Pretty CoreTypeDeclBody where
    pretty :: CoreTypeDeclBody -> Doc AnsiStyle
pretty (CoreTypeAlias Type
t) = Type -> Doc AnsiStyle
prettyTy Type
t
    pretty (CoreDataDecl (Core.TyCon Qualified Text
_ TyConDetails
Core.Prim) [DataCon]
_) = Doc AnsiStyle
"<primitive>"
    pretty (CoreDataDecl TyCon
_ []) = Doc AnsiStyle
"{}"
    pretty (CoreDataDecl TyCon
_ [DataCon]
dcs) = [Doc AnsiStyle] -> Doc AnsiStyle
prettyCtorsInline (DataCon -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty (DataCon -> Doc AnsiStyle) -> [DataCon] -> [Doc AnsiStyle]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DataCon]
dcs)

prettyTdef :: CoreTypeDecl -> Doc AnsiStyle
prettyTdef :: CoreTypeDecl -> Doc AnsiStyle
prettyTdef (CoreTypeDecl Qualified Text
name ElaraKind
kind [TypeVariable]
tvs CoreTypeDeclBody
body) =
    Doc AnsiStyle -> Doc AnsiStyle
keyword Doc AnsiStyle
"type" Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Qualified Text -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty Qualified Text
name Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> [TypeVariable] -> Doc AnsiStyle
prettyTypeVariables [TypeVariable]
tvs Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc AnsiStyle
":" Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ElaraKind -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty ElaraKind
kind Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc AnsiStyle
"=" Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. Doc ann -> Doc ann -> Doc ann
<+> CoreTypeDeclBody -> Doc AnsiStyle
forall a. Pretty a => a -> Doc AnsiStyle
pretty CoreTypeDeclBody
body

makeFields ''CoreModule