{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}

-- | Abstracts traversals over source ASTs.
module Descript.BasicInj.Traverse
  ( module Descript.BasicInj.Traverse.Termed
  , TTerm
  , Traversal (..)
  , Fold (..)
  , Mapping (..)
  , travTerm
  , foldTerm
  , mapTerm
  ) where

import Descript.BasicInj.Traverse.Termed
import Descript.BasicInj.Traverse.Term (TTerm)
import qualified Descript.BasicInj.Traverse.Term as T
import qualified Descript.BasicInj.Data.Value.Reg as Reg
import qualified Descript.BasicInj.Data.Value.In as In
import qualified Descript.BasicInj.Data.Value.Out as Out
import Descript.BasicInj.Data
import Data.Functor.Identity
import Control.Monad ((<=<))
import Control.Monad.Trans.Writer.Strict
import Prelude hiding (mod)

newtype FoldTrav a = FoldTrav a
newtype MappingTrav a = MappingTrav a

class (Monad (Eff f)) => Traversal f where
  type Eff f :: * -> *
  type TAnn f :: *

  tonTerm :: TTerm t -> f -> t (TAnn f) -> Eff f (t (TAnn f))
  tonTerm _ _ = pure

class (Monoid (Res f)) => Fold f where
  type Res f :: *
  type FAnn f :: *

  fonTerm :: TTerm t -> f -> t (FAnn f) -> Res f
  fonTerm _ _ = mempty

class Mapping f where
  type MAnn f :: *

  monTerm :: TTerm t -> f -> t (MAnn f) -> t (MAnn f)
  monTerm _ _ = id

instance (Fold f) => Traversal (FoldTrav f) where
  type Eff (FoldTrav f) = Writer (Res f)
  type TAnn (FoldTrav f) = FAnn f

  tonTerm term (FoldTrav t) x = x <$ tell (fonTerm term t x)

instance (Mapping f) => Traversal (MappingTrav f) where
  type Eff (MappingTrav f) = Identity
  type TAnn (MappingTrav f) = MAnn f

  tonTerm term (MappingTrav t) = Identity . monTerm term t

travTerm :: (Traversal f) => TTerm t -> f -> t (TAnn f) -> Eff f (t (TAnn f))
travTerm term t = travSubTerm term t <=< tonTerm term t

foldTerm :: (Fold f) => TTerm t -> f -> t (FAnn f) -> Res f
foldTerm term t = execWriter . travTerm term (FoldTrav t)

mapTerm :: (Mapping f) => TTerm t -> f -> t (MAnn f) -> t (MAnn f)
mapTerm term t = runIdentity . travTerm term (MappingTrav t)

travSubTerm :: (Traversal f) => TTerm t -> f -> t (TAnn f) -> Eff f (t (TAnn f))
travSubTerm T.Source t (SourceModule mod)
  = SourceModule <$> travTerm T.BModule t mod
travSubTerm T.Source t (SourceProgram prog)
  = SourceProgram <$> travTerm T.Program t prog
travSubTerm T.Program t (Program ann mod query')
  = Program ann <$> travTerm T.BModule t mod <*> travTerm T.Query t query'
travSubTerm T.BModule t (BModule ann ictx amod)
  = BModule ann <$> travTerm T.ImportCtx t ictx <*> travTerm T.AModule t amod
travSubTerm T.AModule t (AModule ann recordCtx' reduceCtx')
    = AModule ann
  <$> travTerm T.RecordCtx t recordCtx'
  <*> travTerm T.ReduceCtx t reduceCtx'
travSubTerm T.ImportCtx t (ImportCtx ann mdecl idecls)
    = ImportCtx ann
  <$> travTerm T.ModuleDecl t mdecl
  <*> traverse (travTerm T.ImportDecl t) idecls
travSubTerm T.RecordCtx t (RecordCtx ann records)
  = RecordCtx ann <$> traverse (travTerm T.RecordDecl t) records
travSubTerm T.ReduceCtx t (ReduceCtx ann topCtx lowCtxs)
    = ReduceCtx ann
  <$> travTerm T.PhaseCtx t topCtx
  <*> traverse (travTerm T.PhaseCtx t) lowCtxs
travSubTerm T.PhaseCtx t (PhaseCtx ann reducers)
  = PhaseCtx ann <$> traverse (travTerm T.Reducer t) reducers
travSubTerm T.Query t (Query ann val) = Query ann <$> travTerm T.RegValue t val
travSubTerm T.ModuleDecl t (ModuleDecl ann path)
  = ModuleDecl ann <$> travTerm T.ModulePath t path
travSubTerm T.ImportDecl t (ImportDecl ann path isrcs idsts)
    = ImportDecl ann
  <$> travTerm T.ModulePath t path
  <*> traverse (travTerm T.ImportRecord t) isrcs
  <*> traverse (travTerm T.ImportRecord t) idsts
travSubTerm T.RecordDecl t (RecordDecl ann recordType)
  = RecordDecl ann <$> travTerm T.RecordType t recordType
travSubTerm T.Reducer t (Reducer ann input' output')
  = Reducer ann <$> travTerm T.Input t input' <*> travTerm T.Output t output'
travSubTerm T.ModulePath t (ModulePath ann xs)
  = ModulePath ann <$> traverse (travTerm T.ModulePathElem t) xs
travSubTerm T.ImportRecord t (ImportRecord ann from to)
    = ImportRecord ann
  <$> travTerm T.RecordHead t from
  <*> travTerm T.RecordHead t to
travSubTerm T.RecordType t (RecordType ann head' properties')
    = RecordType ann
  <$> travTerm T.RecordHead t head'
  <*> traverse (travTerm T.PropertyKey t) properties'
travSubTerm T.RegValue t (Value ann parts)
    = travTerm T.GenValue t
  =<< Value ann <$> traverse (travTerm T.RegPart t) parts
travSubTerm T.Input t (Value ann parts)
    = travTerm T.GenValue t
  =<< Value ann <$> traverse (travTerm T.InPart t) parts
travSubTerm T.Output t (Value ann parts)
    = travTerm T.GenValue t
  =<< Value ann <$> traverse (travTerm T.OutPart t) parts
travSubTerm T.GenValue _ x = pure x
travSubTerm T.RegPart t (Reg.PartPrim prim)
  = Reg.PartPrim <$> travTerm T.Prim t prim
travSubTerm T.RegPart t (Reg.PartRecord record)
  = Reg.PartRecord <$> travTerm T.RegRecord t record
travSubTerm T.InPart t (In.PartPrim prim)
  = In.PartPrim <$> travTerm T.Prim t prim
travSubTerm T.InPart t (In.PartPrimType primType)
  = In.PartPrimType <$> travTerm T.PrimType t primType
travSubTerm T.InPart t (In.PartRecord record)
  = In.PartRecord <$> travTerm T.InRecord t record
travSubTerm T.OutPart t (Out.PartPrim prim)
  = Out.PartPrim <$> travTerm T.Prim t prim
travSubTerm T.OutPart t (Out.PartRecord record)
  = Out.PartRecord <$> travTerm T.OutRecord t record
travSubTerm T.OutPart t (Out.PartPropPath path)
  = Out.PartPropPath <$> travTerm T.PropPath t path
travSubTerm T.OutPart t (Out.PartInjApp app)
  = Out.PartInjApp <$> travTerm T.InjApp t app
travSubTerm T.Prim _ x = pure x
travSubTerm T.PrimType _ x = pure x
travSubTerm T.RegRecord t (Record ann head' properties')
    = travTerm T.GenRecord t
  =<< Record ann head' <$> traverse (travTerm T.RegProperty t) properties'
travSubTerm T.InRecord t (Record ann head' properties')
    = travTerm T.GenRecord t
  =<< Record ann head' <$> traverse (travTerm T.InProperty t) properties'
travSubTerm T.OutRecord t (Record ann head' properties')
    = travTerm T.GenRecord t
  =<< Record ann head' <$> traverse (travTerm T.OutProperty t) properties'
travSubTerm T.GenRecord t (Record ann head' properties')
  = Record ann <$> travTerm T.RecordHead t head' <*> pure properties'
travSubTerm T.RegProperty t (Property ann key val)
    = Property ann
  <$> travTerm T.PropertyKey t key
  <*> travTerm T.RegValue t val
travSubTerm T.InProperty t (Property ann key val)
    = Property ann
  <$> travTerm T.PropertyKey t key
  <*> In.traverseOptVal (travTerm T.Input t) val
travSubTerm T.OutProperty t (Property ann key val)
    = Property ann
  <$> travTerm T.PropertyKey t key
  <*> travTerm T.Output t val
travSubTerm T.GenProperty _ x = pure x
travSubTerm T.PropPath t (PropPath ann elems)
  = PropPath ann <$> traverse (travTerm T.PathElem t) elems
travSubTerm T.PathElem t (PathElem ann propKey' headKey')
    = PathElem ann
  <$> travTerm T.PropertyKey t propKey'
  <*> travTerm T.RecordHead t headKey'
travSubTerm T.InjApp t (Out.InjApp ann funcId' params')
    = Out.InjApp ann funcId'
  <$> traverse (Out.traverseInjParamVal $ travTerm T.Output t) params'
travSubTerm T.ModulePathElem _ x = pure x
travSubTerm T.RecordHead _ x = pure x
travSubTerm T.PropertyKey _ x = pure x