content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  module Kalium.Haskell.Imports where
  
  import Kalium.Prelude hiding (List)
  import Kalium.Util
  import Data.Data
  import qualified Data.Set as S
  import qualified Data.Map as M
  import Language.Haskell.Exts.Syntax
  import Language.Haskell.Exts.SrcLoc (noLoc)
  
  import Control.Monad.Reader
  import Control.Monad.State
  
  imports :: Module -> Module
  imports (Module srcLoc name pragmas wtext exportSpec importDecls decls) =
      let importDecls' = importDecls ++ map importDecl
              (itoList (ModuleName "Prelude" `M.delete` modules))
          modules = gcollect decls
      in unqual (Module srcLoc name pragmas wtext exportSpec importDecls' decls)
          `runReader` mempty
  
  importDecl (moduleName, names) =
      ImportDecl noLoc moduleName
      False False False Nothing Nothing
      (Just (False, IAbs NoNamespace <$> toList names))
  
  gcollect :: Data a => a -> Map ModuleName (Set Name)
  gcollect = (`execState` mempty) . go where
    go :: (Data d, MonadState (Map ModuleName (Set Name)) m)
       => EndoKleisli' m d
    go a | Just (Qual moduleName name) <- cast a
         = a <$ modify (M.insertWith S.union moduleName (S.singleton name))
         | otherwise = gmapM go a
  
  class Unqual a where
      unqual :: (MonadReader (Set Name) m) => EndoKleisli' m a
  
  class    Scoping a    where scoping :: a -> Set Name
  instance Scoping Name where scoping = S.singleton
  
  instance Scoping Pat  where
      scoping = \case
          PWildCard -> mempty
          PVar name -> scoping name
          PTuple _ pats -> scoping pats
          PatTypeSig _ pat _ -> scoping pat
          pat -> error ("unsupported pat: " ++ show pat)
  
  instance Scoping Decl where
      scoping = \case
          FunBind matches -> scoping matches
          PatBind _ pat _ _ -> scoping pat
          TypeSig _ names _ -> scoping names
          decl -> error ("unsupported decl: " ++ show decl)
  
  instance Scoping Stmt where
      scoping = \case
          Qualifier _ -> mempty
          Generator _ pat _ -> scoping pat
          LetStmt binds -> scoping binds
          RecStmt stmts -> scoping stmts
  
  instance Scoping Binds where
      scoping = \case
          IPBinds {} -> error "not supported: IPBinds"
          BDecls decls -> scoping decls
  
  instance Scoping Match where
      scoping (Match _ name _ _ _ _) = scoping name
  
  instance (Scoping a, Foldable f) => Scoping (f a) where
      scoping = foldMap scoping
  
  unqualLocalize
      :: (MonadReader (Set Name) m, Scoping p)
      => p -> Endo' (m a)
  unqualLocalize (scoping -> p) = local (mappend p)
  
  instance (Unqual a, Traversable f) => Unqual (f a) where
      unqual = traverse unqual
  
  instance Unqual Module where
      unqual (Module srcLoc name pragmas wtext exportSpec importDecls            decls)
           =  Module srcLoc name pragmas wtext exportSpec importDecls <$> unqual decls
           & unqualLocalize decls
  
  instance Unqual Decl where
      unqual = \case
          FunBind matches -> FunBind <$> unqual matches
          PatBind srcLoc pat rhs binds
              -> PatBind srcLoc <$> unqual pat <*> unqual rhs <*> unqual binds
               & unqualLocalize pat
          TypeSig srcLoc names ty
              -> TypeSig srcLoc names <$> unqual ty
               & unqualLocalize names
          decl -> error ("unsupported decl: " ++ show decl)
  
  instance Unqual Match where
      unqual (Match srcLoc name pats ty rhs binds)
          = Match srcLoc name
              <$> unqual pats
              <*> unqual ty
              <*> unqual rhs
              <*> unqual binds
  
  instance Unqual Binds where
      unqual = \case
          IPBinds {} -> error "not supported: IPBinds"
          BDecls decls -> BDecls <$> unqual decls
                        & unqualLocalize decls
  
  instance Unqual Rhs where
      unqual = \case
          UnGuardedRhs exp -> UnGuardedRhs <$> unqual exp
          GuardedRhss rhss -> GuardedRhss  <$> unqual rhss
  
  instance Unqual GuardedRhs where
      unqual (GuardedRhs srcLoc            stmts            exp)
           =  GuardedRhs srcLoc <$> unqual stmts <*> unqual exp
           & unqualLocalize stmts
  
  instance Unqual Stmt where
      unqual = \case
          Generator srcLoc pat exp
              -> Generator srcLoc <$> unqual pat <*> unqual exp
               & unqualLocalize pat
          Qualifier exp -> Qualifier <$> unqual exp
          LetStmt binds -> LetStmt <$> unqual binds & unqualLocalize binds
          RecStmt stmts -> RecStmt <$> unqual stmts & unqualLocalize stmts
  
  instance Unqual Exp where
      unqual = \case
          exp@Lit{} -> return exp
          Var qname -> Var <$> unqual qname
          Con qname -> Con <$> unqual qname
          List exps -> List <$> unqual exps
          MultiIf rhss -> MultiIf <$> unqual rhss
          If exp1 exp2 exp3 -> If <$> unqual exp1 <*> unqual exp2 <*> unqual exp3
          EnumFromTo exp1 exp2 -> EnumFromTo <$> unqual exp1 <*> unqual exp2
          Tuple boxed exps -> Tuple boxed <$> unqual exps
          Paren exp -> Paren <$> unqual exp
          App x y -> App <$> unqual x <*> unqual y
          Lambda srcLoc pats exp
               -> Lambda srcLoc
              <$> unqual pats
              <*>(unqual exp & unqualLocalize pats)
          InfixApp x op y -> InfixApp <$> unqual x <*> unqual op <*> unqual y
          RightSection op exp -> RightSection <$> unqual op <*> unqual exp
          Do stmts -> Do <$> unqual stmts & unqualLocalize stmts
          Let binds exp
              -> Let <$> unqual binds <*> unqual exp
               & unqualLocalize binds
          exp -> error ("unsupported exp: " ++ show exp)
  
  instance Unqual Type where
      unqual = \case
          TyFun ty1 ty2 -> TyFun <$> unqual ty1 <*> unqual ty2
          TyApp ty1 ty2 -> TyApp <$> unqual ty1 <*> unqual ty2
          TyTuple boxed ty -> TyTuple boxed <$> unqual ty
          TyList ty -> TyList <$> unqual ty
          TyCon qname -> TyCon <$> unqual qname
          ty -> error ("unsupported type: " ++ show ty)
  
  instance Unqual Pat where
      unqual = \case
          pat@PWildCard -> return pat
          pat@PVar{}    -> return pat
          PTuple boxed pats -> PTuple boxed <$> unqual pats
          PatTypeSig srcLoc pat ty ->
              PatTypeSig srcLoc <$> unqual pat <*> unqual ty
          pat -> error ("unsupported pat: " ++ show pat)
  
  instance Unqual QOp where
      unqual = \case
          QVarOp qname -> QVarOp <$> unqual qname
          QConOp qname -> QConOp <$> unqual qname
  
  instance Unqual QName where
      unqual qname@(Qual _ name) =
          asks (elem name) <&> bool (UnQual name) qname
      unqual qname = return qname
  
  