content: module Vehicle.Syntax.AST.Decl where
  
  import Control.DeepSeq (NFData)
  import Data.Hashable (Hashable)
  import Data.Serialize (Serialize)
  import GHC.Generics (Generic)
  import Prettyprinter (Pretty (..))
  import Vehicle.Syntax.AST.Name
  import Vehicle.Syntax.AST.Provenance
  import Vehicle.Syntax.AST.Record
  import Vehicle.Syntax.AST.Type
  
  --------------------------------------------------------------------------------
  -- Declarations
  
  -- | Type of top-level declarations.
  data GenericDecl expr
    = -- | Abstract definitions that require no body
      DefAbstract
        Provenance -- Location in source file.
        Identifier -- Name of definition.
        DefAbstractSort -- The sort of abstract definition.
        expr -- Type of the definition.
    | -- | Function definitions with a body
      DefFunction
        Provenance -- Location in source file.
        Identifier -- Name of definition.
        [Annotation] -- List of annotations.
        expr -- Type of the definition.
        expr -- Body of the definition.
    | -- | Record definitions
      DefRecord
        Provenance -- Location in source file.
        Identifier -- Name of definition.
        [Annotation] -- List of annotations.
        expr -- Type of the record
        (RecordFields expr) -- Fields in the record definition.
    deriving (Show, Functor, Foldable, Traversable, Generic)
  
  instance (NFData expr) => NFData (GenericDecl expr)
  
  instance (Serialize expr) => Serialize (GenericDecl expr)
  
  instance HasProvenance (GenericDecl expr) where
    provenanceOf = \case
      DefAbstract p _ _ _ -> p
      DefFunction p _ _ _ _ -> p
      DefRecord p _ _ _ _ -> p
  
  instance HasIdentifier (GenericDecl expr) where
    identifierOf = \case
      DefAbstract _ i _ _ -> i
      DefFunction _ i _ _ _ -> i
      DefRecord _ i _ _ _ -> i
  
  instance HasName (GenericDecl expr) Name where
    nameOf = nameOf . identifierOf
  
  instance HasType (GenericDecl expr) expr where
    typeOf = \case
      DefAbstract _ _ _ t -> t
      DefFunction _ _ _ t _ -> t
      DefRecord _ _ _ t _ -> t
  
  bodyOf :: GenericDecl expr -> Maybe expr
  bodyOf = \case
    DefFunction _ _ _ _ e -> Just e
    DefAbstract {} -> Nothing
    DefRecord {} -> Nothing
  
  annotationsOf :: GenericDecl expr -> [Annotation]
  annotationsOf = \case
    DefFunction _ _ anns _ _ -> anns
    DefAbstract {} -> []
    DefRecord _ _ anns _ _ -> anns
  
  abstractSortOf :: GenericDecl expr -> Maybe DefAbstractSort
  abstractSortOf decl = case decl of
    DefAbstract _ _ sort _ -> Just sort
    DefFunction {} -> Nothing
    DefRecord {} -> Nothing
  
  -- | Traverses the type and body of a declaration using the first and
  -- second provided functions respectively.
  -- Use |traverse| if you want to traverse them using the same function.
  traverseDeclTypeAndExpr ::
    (Monad m) =>
    (expr1 -> m expr2) ->
    (expr1 -> m expr2) ->
    GenericDecl expr1 ->
    m (GenericDecl expr2)
  traverseDeclTypeAndExpr f1 f2 = \case
    DefAbstract p n r t -> DefAbstract p n r <$> f1 t
    DefFunction p n b t e -> DefFunction p n b <$> f1 t <*> f2 e
    DefRecord p n b t fs -> DefRecord p n b <$> f1 t <*> traverseRecordFields f2 fs
  
  mapIdentifier ::
    (Identifier -> Identifier) ->
    GenericDecl expr ->
    GenericDecl expr
  mapIdentifier f = \case
    DefAbstract p n r t -> DefAbstract p (f n) r t
    DefFunction p n b t e -> DefFunction p (f n) b t e
    DefRecord p n b t fs -> DefRecord p (f n) b t fs
  
  -- | Traverses the type of the declaration.
  traverseDeclType ::
    (Monad m) =>
    (expr -> m expr) ->
    GenericDecl expr ->
    m (GenericDecl expr)
  traverseDeclType f = traverseDeclTypeAndExpr f return
  
  isPropertyDecl :: GenericDecl expr -> Bool
  isPropertyDecl = \case
    DefAbstract {} -> False
    DefFunction _ _ anns _ _ -> AnnProperty `elem` anns
    DefRecord {} -> False
  
  isAbstractDecl :: GenericDecl expr -> Bool
  isAbstractDecl = \case
    DefAbstract {} -> True
    DefFunction {} -> False
    DefRecord {} -> False
  
  --------------------------------------------------------------------------------
  -- Abstract definition types options
  
  data ParameterSort
    = Inferable
    | NonInferable
    deriving (Eq, Ord, Show, Generic)
  
  instance NFData ParameterSort
  
  instance Serialize ParameterSort
  
  instance Hashable ParameterSort
  
  instance Pretty ParameterSort where
    pretty = \case
      Inferable -> "(infer=True)"
      NonInferable -> ""
  
  isInferable :: ParameterSort -> Bool
  isInferable = \case
    Inferable -> True
    NonInferable -> False
  
  data DefAbstractSort
    = NetworkDef
    | DatasetDef
    | ParameterDef ParameterSort
    | PostulateDef
    deriving (Eq, Show, Generic)
  
  instance NFData DefAbstractSort
  
  instance Serialize DefAbstractSort
  
  instance Pretty DefAbstractSort where
    pretty t =
      "@" <> case t of
        NetworkDef -> "network"
        DatasetDef -> "dataset"
        ParameterDef {} -> "parameter"
        PostulateDef {} -> "postulate"
  
  isExternalResourceSort :: DefAbstractSort -> Bool
  isExternalResourceSort = \case
    NetworkDef -> True
    DatasetDef -> True
    ParameterDef parameterType -> parameterType == NonInferable
    PostulateDef {} -> False
  
  isExternalResourceDecl :: GenericDecl expr -> Bool
  isExternalResourceDecl decl = maybe False isExternalResourceSort (abstractSortOf decl)
  
  convertToPostulate :: GenericDecl expr -> GenericDecl expr
  convertToPostulate d =
    DefAbstract (provenanceOf d) (identifierOf d) PostulateDef (typeOf d)
  
  --------------------------------------------------------------------------------
  -- Annotations options
  
  data Annotation
    = AnnProperty
    | AnnTensor
    deriving (Eq, Show, Generic)
  
  instance NFData Annotation
  
  instance Serialize Annotation
  
  instance Pretty Annotation where
    pretty = \case
      AnnProperty -> "@property"
      AnnTensor -> "@tensor"
  
  isAnnotatedAsProperty :: [Annotation] -> Bool
  isAnnotatedAsProperty anns = AnnProperty `elem` anns
  
  isAnnotatedAsTensor :: [Annotation] -> Bool
  isAnnotatedAsTensor anns = AnnTensor `elem` anns
  