{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}

module AST.Structure
    ( FixAST, ASTNS, ASTNS1
    , foldReferences
    , bottomUpReferences
    , mapNs
    ) where


import Data.Coapplicative
import Data.Foldable (fold)
import Data.Functor.Const
import AST.V0_16
import qualified Data.Indexed as I


-- FixAST :: (Type -> Type) -> Type -> Type -> Type -> NodeKind -> Type
type FixAST annf typeRef ctorRef varRef =
    I.Fix annf (AST typeRef ctorRef varRef)

-- ASTNS :: (Type -> Type) -> Type -> NodeKind -> Type
type ASTNS annf ns =
    FixAST annf (ns, UppercaseIdentifier) (ns, UppercaseIdentifier) (Ref ns)

-- This is the same as ASTNS, but with the first level unFix'ed
-- ASTNS1 :: (Type -> Type) -> Type -> NodeKind -> Type
type ASTNS1 annf ns =
    AST
        (ns, UppercaseIdentifier)
        (ns, UppercaseIdentifier)
        (Ref ns)
        (ASTNS annf ns)


bottomUpReferences ::
    (Functor annf) =>
    (typeRef1 -> typeRef2)
    -> (ctorRef1 -> ctorRef2)
    -> (varRef1 -> varRef2)
    -> (forall kind.
        FixAST annf typeRef1 ctorRef1 varRef1 kind
        -> FixAST annf typeRef2 ctorRef2 varRef2 kind
       )
bottomUpReferences ftr fcr fvr =
    I.cata (I.Fix . fmap (mapAll ftr fcr fvr id))


foldReferences ::
    forall a annf typeRef ctorRef varRef kind.
    (Monoid a, Coapplicative annf) =>
    (typeRef -> a) -> (ctorRef -> a) -> (varRef -> a)
    -> FixAST annf typeRef ctorRef varRef kind -> a
foldReferences ftype fctor fvar =
    getConst . I.cata (foldNode  . extract)
    where
        -- This is kinda confusing, but we use the Const type constructor to merge all the different NodeKinds into a single type `a`
        -- See http://www.timphilipwilliams.com/posts/2013-01-16-fixing-gadts.html for relevant details.
        foldNode :: forall kind'. AST typeRef ctorRef varRef (Const a) kind' -> Const a kind'
        foldNode = \case
            TopLevel tls -> Const $ foldMap (foldMap getConst) tls

            -- Declarations
            Definition name args _ e -> Const (getConst name <> foldMap (getConst . extract) args <> getConst e)
            TypeAnnotation _ t -> Const (getConst $ extract t)
            CommonDeclaration d -> Const $ getConst d
            Datatype _ ctors -> Const (foldMap (getConst . fold) ctors)
            TypeAlias _ _ t -> Const (getConst $ extract t)
            PortAnnotation _ _ t -> Const (getConst t)
            PortDefinition_until_0_16 _ _ e -> Const (getConst e)
            Fixity_until_0_18 _ _ _ _ name -> Const (fvar name)
            Fixity _ _ _ _ -> mempty

            -- Expressions
            Unit _ -> mempty
            Literal _ -> mempty
            VarExpr var -> Const $ fvar var
            App first rest _ -> first <> mconcat (fmap extract rest)
            Unary _ e -> e
            Binops first ops _ -> Const (getConst first <> foldMap foldBinopsClause ops)
            Parens e -> extract e
            ExplicitList terms _ _ -> fold terms
            Range left right _ -> extract left <> extract right
            Tuple terms _ -> mconcat $ fmap extract terms
            TupleFunction _ -> mempty
            Record _ fields _ _ -> foldMap (extract . _value) fields
            Access e _ -> e
            AccessFunction _ -> mempty
            Lambda args _ e _ -> Const (foldMap (getConst . extract) args <> getConst e)
            If cond elsifs els -> Const (foldIfClause cond <> foldMap (foldIfClause . extract) elsifs <> getConst (extract els))
            Let defs _ e -> Const (foldMap getConst defs <> getConst e)
            LetCommonDeclaration d -> Const $ getConst d
            LetComment _ -> mempty
            Case (cond, _) branches -> Const (getConst (extract cond) <> foldMap getConst branches)
            CaseBranch _ _ _ p e -> Const (getConst p <> getConst e)
            GLShader _ -> mempty

            -- Patterns
            Anything -> mempty
            UnitPattern _ -> mempty
            LiteralPattern _ -> mempty
            VarPattern _ -> mempty
            OpPattern _ -> mempty
            DataPattern ctor args -> Const (fctor ctor <> foldMap (getConst . extract) args)
            PatternParens p -> extract p
            TuplePattern terms -> foldMap extract terms
            EmptyListPattern _ -> mempty
            ListPattern terms -> foldMap extract terms
            ConsPattern first rest -> extract first <> fold rest
            EmptyRecordPattern _ -> mempty
            RecordPattern _ -> mempty
            Alias p _ -> extract p

            -- Types
            UnitType _ -> mempty
            TypeVariable _ -> mempty
            TypeConstruction name args _ -> Const (foldTypeConstructor name <> foldMap (getConst . extract) args)
            TypeParens typ -> extract typ
            TupleType terms _ -> foldMap extract terms
            RecordType _ fields _ _ -> foldMap (extract . _value) fields
            FunctionType first rest _ -> extract first <> fold rest

        foldTypeConstructor :: TypeConstructor typeRef -> a
        foldTypeConstructor = \case
            NamedConstructor name -> ftype name
            TupleConstructor _ -> mempty

        foldBinopsClause :: BinopsClause varRef (Const a 'ExpressionNK) -> a
        foldBinopsClause = \case
            BinopsClause _ op _ e -> fvar op <> getConst e

        foldIfClause :: IfClause (Const a 'ExpressionNK) -> a
        foldIfClause = \case
            IfClause cond els -> getConst (extract cond) <> getConst (extract els)


mapNs ::
    Functor annf =>
    (ns1 -> ns2)
    -> (forall kind.
        ASTNS annf ns1 kind
        -> ASTNS annf ns2 kind
       )
mapNs f =
    let
        mapTypeRef (ns, u) = (f ns, u)
        mapCtorRef (ns, u) = (f ns, u)
        mapVarRef (VarRef ns l) = VarRef (f ns) l
        mapVarRef (TagRef ns u) = TagRef (f ns) u
        mapVarRef (OpRef op) = OpRef op
    in
    bottomUpReferences mapTypeRef mapCtorRef mapVarRef
