content: {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  {-# LANGUAGE OverloadedStrings #-}
  module Icicle.Compiler.Source
    ( ErrorSource (..)
  
    , IsName
  
    , Var
    , TypeAnnot
    , AnnotUnit
  
    , QueryUntyped
    , QueryTyped
  
    , CoreProgramUntyped
  
    , Funs
    , FunEnvT
  
      -- * Compiler options
    , IcicleCompileOptions (..)
    , defaultCompileOptions
    , FusionOptions (..)
    , defaultFusionOptions
    , Inline.InlineOption (..)
    , Inline.defaultInline
    , Check.CheckOptions (..)
    , Check.defaultCheckOptions
  
      -- * From dictionaires and libraries
    , queryOfSource
    , entryOfQuery
    , readIcicleLibrary
  
      -- * Works on Source programs
    , sourceParseQT
    , sourceParseF
    , sourceDesugarQT
    , sourceDesugarF
    , sourceReifyQT
    , sourceCheckQT
    , sourceCheckF
    , sourceCheckFunLog
    , sourceInline
  
      -- * Helpers
    , freshNamer
    , annotOfError
    ) where
  
  
  import           Icicle.Common.Base                       (Name)
  import qualified Icicle.Common.Fresh                      as Fresh
  
  import qualified Icicle.Core.Program.Program              as Core
  
  import           Icicle.Data.Name
  
  import           Icicle.Dictionary                        (Dictionary, DictionaryOutput(..))
  import qualified Icicle.Dictionary                        as Dict
  
  import           Icicle.Internal.Pretty
  
  import qualified Icicle.Source.Checker                    as Check
  import qualified Icicle.Source.Parser                     as Parse
  import qualified Icicle.Source.Query                      as Query
  import qualified Icicle.Source.Transform.Desugar          as Desugar
  import qualified Icicle.Source.Transform.Inline           as Inline
  import qualified Icicle.Source.Transform.ReifyPossibility as Reify
  import qualified Icicle.Source.Type                       as Type
  
  import           Data.Functor.Identity
  import qualified Data.Map                                 as M
  import           Data.String
  import           Data.Hashable                            (Hashable)
  
  import qualified Text.ParserCombinators.Parsec            as Parsec
  
  import           GHC.Generics                             (Generic)
  
  import           P
  
  import           X.Control.Monad.Trans.Either
  
  
  
  type IsName v = (Hashable v, Eq v, IsString v, Pretty v, Show v, NFData v)
  
  type AnnotUnit = ()
  
  type Var         = Parse.Variable
  type TypeAnnot a = Type.Annot a Var
  
  type Funs a b = [((a, Name b), Query.Function a b)]
  type FunEnvT a b = [ ( Name b
                     , ( Type.FunctionType b
                       , Query.Function (TypeAnnot a) b )) ]
  
  type QueryUntyped v = Query.QueryTop            Parsec.SourcePos  v
  type QueryTyped   v = Query.QueryTop (TypeAnnot Parsec.SourcePos) v
  
  type CoreProgramUntyped v = Core.Program AnnotUnit v
  
  type Error = ErrorSource Var
  
  freshNamer :: IsString v => v -> Fresh.NameState v
  freshNamer prefix
   = Fresh.counterPrefixNameState (fromString . show) prefix
  
  --------------------------------------------------------------------------------
  
  data IcicleCompileOptions = IcicleCompileOptions
    { icicleInline        :: Inline.InlineOption
    , icicleBigData       :: Check.CheckOptions
    , icicleFusionOptions :: FusionOptions
    }
  
  defaultCompileOptions :: IcicleCompileOptions
  defaultCompileOptions = IcicleCompileOptions Inline.defaultInline Check.defaultCheckOptions defaultFusionOptions
  
  data FusionOptions = FusionOptions
    { fusionMaximumPerKernel :: Int
    }
  
  defaultFusionOptions :: FusionOptions
  defaultFusionOptions = FusionOptions 100
  
  --------------------------------------------------------------------------------
  
  data ErrorSource var
   = ErrorSourceParse       !Parsec.ParseError
   | ErrorSourceDesugar     !(Desugar.DesugarError Parsec.SourcePos var)
   | ErrorSourceCheck       !(Check.CheckError     Parsec.SourcePos var)
   | ErrorSourceResolveError !UnresolvedInputId
   deriving (Show, Generic)
  
  -- FIXME We can't implement NFData properly for this type because Parsec.SourcePos is
  -- FIXME not NFData, we really should define our own type for source positions.
  instance NFData (ErrorSource a) where rnf x = seq x ()
  
  annotOfError :: ErrorSource a -> Maybe Parsec.SourcePos
  annotOfError e
   = case e of
      ErrorSourceParse sp
       -> Just
        $ Parsec.errorPos sp
      ErrorSourceDesugar e'
       -> Desugar.annotOfError e'
      ErrorSourceCheck       e'
       -> Check.annotOfError  e'
      ErrorSourceResolveError _
       -> Nothing
  
  instance (Hashable a, Eq a, IsString a, Pretty a) => Pretty (ErrorSource a) where
    pretty = \case
      ErrorSourceParse p ->
        vsep [
            reannotate AnnErrorHeading $ prettyH2 "Parse error"
          , mempty
          , indent 2 . text $ show p
          ]
  
      ErrorSourceDesugar d ->
        vsep [
            reannotate AnnErrorHeading $ prettyH2 "Desugar error"
          , mempty
          , indent 2 $ pretty d
          ]
  
      ErrorSourceCheck ce ->
        vsep [
            reannotate AnnErrorHeading $ prettyH2 "Check error"
          , mempty
          , indent 2 $ pretty ce
          ]
  
      ErrorSourceResolveError t ->
        vsep [
            reannotate AnnErrorHeading $ prettyH2 "Could not resolve input name"
          , mempty
          , indent 2 $ pretty t
          ]
  
  --------------------------------------------------------------------------------
  
  -- * queries
  
  queryOfSource :: Check.CheckOptions
                -> Dictionary
                -> OutputId
                -> Text
                -> Either Error (QueryTyped Var)
  queryOfSource checkOpts dict oid src = do
    parsed       <- sourceParseQT oid src
    desugared    <- sourceDesugarQT parsed
    (checked, _) <- sourceCheckQT checkOpts dict desugared
    pure checked
  
  entryOfQuery ::
       Namespace
    -> OutputName
    -> QueryTyped Var
    -> DictionaryOutput
  entryOfQuery nsp oname query
    = DictionaryOutput (OutputId nsp oname) query
  
  -- * source
  
  sourceParseQT :: OutputId
                -> Text
                -> Either Error (QueryUntyped Var)
  sourceParseQT oid t
   = first ErrorSourceParse
   $ Parse.parseQueryTop oid t
  
  sourceParseF :: Parsec.SourceName
               -> Text
               -> Either Error (Funs Parsec.SourcePos Var)
  sourceParseF env t
   = first ErrorSourceParse
   $ Parse.parseFunctions env t
  
  sourceDesugarQT ::               QueryUntyped Var
                  -> Either Error (QueryUntyped Var)
  sourceDesugarQT q
   = runIdentity . runEitherT . bimapEitherT ErrorSourceDesugar snd
   $ Fresh.runFreshT
       (Desugar.desugarQT q)
       (freshNamer "desugar_q")
  
  sourceDesugarF :: Funs Parsec.SourcePos Var
                 -> Either (ErrorSource Var) (Funs Parsec.SourcePos Var)
  sourceDesugarF fun
   = runIdentity . runEitherT . bimapEitherT ErrorSourceDesugar snd
   $ Fresh.runFreshT
       (mapM (mapM Desugar.desugarFun) fun)
       (freshNamer "desugar_f")
  
  sourceReifyQT :: QueryTyped Var
                -> QueryTyped Var
  sourceReifyQT q
   = snd
   $ runIdentity
   $ Fresh.runFreshT
       (Reify.reifyPossibilityQT q)
       (freshNamer "reify")
  
  sourceCheckQT :: Check.CheckOptions
                -> Dictionary
                -> QueryUntyped Var
                -> Either Error (QueryTyped Var, Type.Type Var)
  sourceCheckQT opts d q
   = let d' = Dict.featureMapOfDictionary d
     in  first ErrorSourceCheck
       $ snd
       $ flip Fresh.runFresh (freshNamer "check")
       $ runEitherT
       $ Check.checkQT opts d' q
  
  sourceCheckF :: FunEnvT Parsec.SourcePos Var
               -> Funs    Parsec.SourcePos Var
               -> Either  Error (FunEnvT Parsec.SourcePos Var)
  sourceCheckF env parsedImport
   = second fst
   $ sourceCheckFunLog env parsedImport
  
  sourceCheckFunLog :: FunEnvT Parsec.SourcePos Var
               -> Funs    Parsec.SourcePos Var
               -> Either  Error (FunEnvT Parsec.SourcePos Var, [[Check.CheckLog Parsec.SourcePos Var]])
  sourceCheckFunLog env parsedImport
   = first ErrorSourceCheck
   $ snd
   $ flip Fresh.runFresh (freshNamer "check")
   $ runEitherT
   $ Check.checkFs env parsedImport
  
  sourceInline :: Inline.InlineOption
               -> Dictionary
               -> QueryTyped   Var
               -> QueryUntyped Var
  sourceInline opt d q
   = Query.reannotQT Type.annAnnot
   $ inline q
   where
    funs      = M.fromList
              $ fmap (\x -> (Dict.functionName x, Dict.functionDefinition x))
              $ Dict.dictionaryFunctions d
    inline q' = snd
              $ Fresh.runFresh
                  (Inline.inlineQT opt funs q')
                  (freshNamer "inline")
  
  readIcicleLibrary :: Var -> Parsec.SourceName -> Text -> Either Error (FunEnvT Parsec.SourcePos Var)
  readIcicleLibrary name source input
   = do input' <- first ErrorSourceParse $ Parse.parseFunctions source input
        first ErrorSourceCheck
               $ second fst
               $ snd
               $ flip Fresh.runFresh (freshNamer name)
               $ runEitherT
               $ Check.checkFs [] input'
  