-----------------------------------------------------------------------------
-- |The Helium Compiler : Static Analysis
-- 
-- Maintainer  :  bastiaan@cs.uu.nl
-- Stability   :  experimental
-- Portability :  unknown
--
-- Kind inference.
--
---------------------------------------------------------------------------------------

-------------------------------------------------------------------------------
-- Attributes for kind inferencing

ATTR Module
   [ importEnvironment : ImportEnvironment
     options           : {[Option]}
   |
   | debugIO           : {IO ()}
     kindEnvironment   : KindEnvironment
     kindErrors        : KindErrors
   ]
   
-------------------------------------------------------------------------------
-- Includes

INCLUDE "UHA_Syntax.ag"

-------------------------------------------------------------------------------
-- Imports

imports {
import Top.Types
import Top.Solver.Greedy
import Top.Solver
import Helium.StaticAnalysis.Miscellaneous.TypeConstraints
import Helium.Syntax.UHA_Syntax
import Helium.Main.Args
import qualified Data.Map as M

-- import StaticAnalysis.Miscellaneous.TypeConstraints
import Helium.Utils.Utils (internalError)
import Helium.ModuleSystem.ImportEnvironment hiding (setTypeSynonyms)
import Helium.StaticAnalysis.Messages.KindErrors
import Data.Char (isLower)
import Helium.StaticAnalysis.Inferencers.BindingGroupAnalysis (Assumptions, PatternAssumptions, noAssumptions, combine, single, topSort) 

}

-------------------------------------------------------------------------------
-- Semantic functions

SEM Module
  | Module   lhs  . kindErrors      = @substitution |-> (map fst @kindErrors)
                  . debugIO         = putStrLn (show @logEntries)
             body . kappaUnique     = 0
             loc  . (SolveResult kappaUniqueAtTheEnd substitution _ _ kindErrors, logEntries) = 
                       solve (solveOptions { uniqueCounter = @body.kappaUnique }) @body.constraints greedyConstraintSolver
                  . kindEnvironment = let f kind = generalizeAll ([] .=>. defaultToStar (@substitution |-> kind))
                                      in M.map f @body.environment                     
                                        

ATTR Body [ importEnvironment:ImportEnvironment | | constraints:KindConstraints environment:PatternAssumptions ]

SEM Body 
  | Hole
       lhs . constraints = []
           . environment = noAssumptions
  | Body 
       declarations . bindingGroups  = []
       lhs . constraints             = @newConstraints ++ @cs
       loc . (environment, aset, cs) = performBindingGroup @declarations.bindingGroups
           . newConstraints          = fst $ (@kindEnvironment .:::. @aset) (\n -> unexpected $ "Body.Body " ++ show n)
           . kindEnvironment         = getKindsFromImportEnvironment @lhs.importEnvironment

ATTR Body Declarations Declaration Expressions Expression Alternatives Alternative
     Statement Statements Qualifier Qualifiers RightHandSide FunctionBinding
     FunctionBindings Constructor Constructors RecordExpressionBinding RecordExpressionBindings
     MaybeExpression MaybeDeclarations GuardedExpression GuardedExpressions
     FieldDeclaration FieldDeclarations ContextItems ContextItem
     AnnotatedType AnnotatedTypes Type Types SimpleType
     [ | kappaUnique:Int | ]

ATTR Declaration Declarations Expressions Expression Alternatives Alternative 
     Statement Statements Qualifier Qualifiers RightHandSide FunctionBinding
     FunctionBindings MaybeExpression RecordExpressionBinding RecordExpressionBindings
     MaybeDeclarations GuardedExpression GuardedExpressions
     [ | bindingGroups:BindingGroups | ]

SEM Declaration
  | Type
       simpletype . constraints = []
                  . kappaOfRHS  = @type.kappa
       lhs  . bindingGroups  = @newGroup : @lhs.bindingGroups
       loc  . newConstraints = fst $ (@simpletype.environment .===. @type.assumptions) (\n -> unexpected $ "Declaration.Type " ++ show n)
            . newGroup       = (@simpletype.declared, @type.assumptions, @newConstraints ++ @type.constraints)
            
  | Data
       simpletype . constraints = [] 
                  . kappaOfRHS  = star
       lhs . bindingGroups  = @newGroup : @lhs.bindingGroups
       loc . newConstraints = fst $ (@simpletype.environment .===. @constructors.assumptions) (\n -> unexpected $ "Declaration.Data " ++ show n)
           . newGroup       = (@simpletype.declared, @constructors.assumptions, @newConstraints ++ @constructors.constraints)
           
  | TypeSignature
       type . constraints  = [] 
       lhs  . bindingGroups = @newGroup : @lhs.bindingGroups
            . kappaUnique   = @type.kappaUnique + length @tvEnv
       loc  . newConstraint = (@type.kappa <==> star) (mustBeStar @range.self "type signature" @type.self)
            . tvEnv         = zip (getTypeVariables @type.assumptions) (map TVar [@type.kappaUnique..])
            . (cset, aset)  = (M.fromList @tvEnv .===. @type.assumptions) (\n -> unexpected $ "Declaration.TypeSignature " ++ show n) 
            . newGroup      = (M.empty, @aset, @cset ++ @type.constraints ++ [@newConstraint])


SEM Expression
  | Typed
       type . constraints   = []   
       lhs  . bindingGroups = @newGroup : @expression.bindingGroups
            . kappaUnique   = @type.kappaUnique + length @tvEnv
       loc  . newConstraint = (@type.kappa <==> star) (mustBeStar @range.self "type annotation" @type.self)
            . tvEnv         = zip (getTypeVariables @type.assumptions) (map TVar [@type.kappaUnique..])
            . (cset, aset)  = (M.fromList @tvEnv .===. @type.assumptions) (\n -> unexpected $ "Expression.Typed " ++ show n) 
            . newGroup      = (M.empty, @aset, @cset ++ @type.constraints ++ [@newConstraint])

ATTR SimpleType [ kappaOfRHS:Kind | constraints:KindConstraints | environment:PatternAssumptions declared:PatternAssumptions ]

SEM SimpleType
  | SimpleType
       lhs . environment   = M.fromList (zip @typevariables.self @kappasVars)
           . declared      = M.singleton @name.self @kappaCon
           . constraints   = @newConstraint : @lhs.constraints
           . kappaUnique   = 1 + length @typevariables.self + @lhs.kappaUnique
       loc . kappaCon      = TVar @lhs.kappaUnique
           . kappasVars    = take (length @typevariables.self) [ TVar i | i <- [ @lhs.kappaUnique+1 .. ]]
           . newConstraint = (@kappaCon .==. foldr (.->.) @lhs.kappaOfRHS @kappasVars) (unexpected "SimpleType.SimpleType")

ATTR AnnotatedTypes Types [ | constraints:KindConstraints | assumptions:Assumptions kappas:Kinds]
ATTR AnnotatedType  Type  [ | constraints:KindConstraints | assumptions:Assumptions kappa:Kind ]

SEM Type
  | Application
       lhs . assumptions   = @function.assumptions `combine` @arguments.assumptions
           . constraints   = @arguments.constraints ++ [@newConstraint]
       lhs . kappaUnique   = @arguments.kappaUnique + 1
       loc . kappa         = TVar @arguments.kappaUnique
           . newConstraint = (@function.kappa <==> foldr (.->.) @kappa @arguments.kappas) (kindApplication @range.self @self @function.self)
           
           
  | Variable 
       lhs . assumptions = single @name.self @kappa           
           . kappaUnique = @lhs.kappaUnique + 1
       loc . kappa       = TVar @lhs.kappaUnique 
                 
  | Constructor   
       lhs . assumptions = single @name.self @kappa
           . kappaUnique = @lhs.kappaUnique + 1
       loc . kappa       = TVar @lhs.kappaUnique         

SEM Types 
  | Cons   lhs . assumptions = @hd.assumptions `combine` @tl.assumptions
               . kappas      = @hd.kappa : @tl.kappas   
  | Nil    lhs . assumptions = noAssumptions
               . kappas      = []  

SEM AnnotatedType 
  | AnnotatedType
       lhs . constraints   = @type.constraints ++ [@newConstraint]
       loc . newConstraint = (@type.kappa <==> star) (mustBeStar @range.self "data type declaration" @type.self)

SEM AnnotatedTypes
  | Cons   lhs . assumptions = @hd.assumptions `combine` @tl.assumptions
               . kappas      = @hd.kappa : @tl.kappas
  | Nil    lhs . assumptions = noAssumptions
               . kappas      = []

ATTR Constructor Constructors [ | constraints:KindConstraints | assumptions:Assumptions ]

SEM Constructor
  | Infix   lhs . assumptions = @leftType.assumptions `combine` @rightType.assumptions

SEM Constructors
  | Cons   lhs . assumptions = @hd.assumptions `combine` @tl.assumptions
  | Nil    lhs . assumptions = noAssumptions
  

SEM FieldDeclaration 
  | FieldDeclaration    loc . constraints = internalError "KindInferencing.ag" "n/a" "Field decls are not supported"
SEM Declaration 
  | Instance loc . constraints = internalError "KindInferencing.ag" "n/a" "instance decls are not supported"
  | Default  loc . constraints = internalError "KindInferencing.ag" "n/a" "default decls is not supported"
  | Class    loc . (constraints,kappaOfRHS) = internalError "KindInferencing.ag" "n/a" "class decls are not supported"
  | Newtype  loc . (constraints,kappaOfRHS) = internalError "KindInferencing.ag" "n/a" "newtype decls are not supported"
SEM ContextItem 
  | ContextItem    loc . constraints = internalError "KindInferencing.ag" "n/a" "ContextItems are not supported"
SEM Type  
  | Qualified   loc . (assumptions, kappa) = internalError "KindInferencing.ag" "n/a" "Qualified types are not supported"
  | Forall      loc . (assumptions, kappa) = internalError "KindInferencing.ag" "n/a" "Universal types are not supported"
  | Exists      loc . (assumptions, kappa) = internalError "KindInferencing.ag" "n/a" "Existential types are not supported"
SEM Constructor
  | Record  lhs . assumptions = internalError "KindInferencing.ag" "n/a" "Record constructors are not supported"

{

type KindEnvironment = M.Map Name TpScheme
type KindConstraint  = TypeConstraint  KindError
type KindConstraints = TypeConstraints KindError

type BindingGroups = [BindingGroup]
type BindingGroup  = (PatternAssumptions,Assumptions,KindConstraints)

combineBindingGroup :: BindingGroup -> BindingGroup -> BindingGroup
combineBindingGroup (e1,a1,c1) (e2,a2,c2) = (e1 `M.union` e2,a1 `combine` a2,c1++c2)

concatBindingGroups :: BindingGroups -> BindingGroup
concatBindingGroups = foldr combineBindingGroup emptyBindingGroup

emptyBindingGroup :: BindingGroup
emptyBindingGroup = (noAssumptions, noAssumptions, [])

performBindingGroup :: BindingGroups -> (PatternAssumptions, Assumptions, KindConstraints)
performBindingGroup = glueGroups . bindingGroupAnalysis    
   where   
      bindingGroupAnalysis :: BindingGroups -> BindingGroups
      bindingGroupAnalysis cs
         = let indexMap = concat (zipWith f cs [0..])
               f (env,_,_) i = [ (n,i) | n <- M.keys env ]
               edges    = concat (zipWith f' cs [0..])
               f' (_,ass,_) i = [ (i,j)| n <- M.keys ass, (n',j) <- indexMap, n==n' ]
               list = topSort (length cs-1) edges
           in map (concatBindingGroups . map (cs !!)) list
           
      glueGroups :: BindingGroups -> (PatternAssumptions, Assumptions, KindConstraints)
      glueGroups = foldr op (noAssumptions, noAssumptions, []) 
         where
            op (env, aset, cset) (environment, assumptions, constraints) = 
               let (cset1,aset')        = (env .===. aset)            (\n -> unexpected $ "BindingGroup.same "++show n)
                   (cset2,assumptions') = (!<==!) [] env assumptions  (\n -> unexpected $ "BindingGroup.instance "++show n)
               in ( env `M.union` environment
                  , aset' `combine` assumptions'
                  , cset1 ++ cset ++ cset2 ++ constraints
                  )             

getKindsFromImportEnvironment :: ImportEnvironment -> KindEnvironment
getKindsFromImportEnvironment = M.map f . typeConstructors
   where f i = generalizeAll ([] .=>. foldr (.->.) star (replicate i star))

getTypeVariables :: Assumptions -> Names
getTypeVariables = filter p . M.keys
   where p n = case show n of
                  []  -> False
                  c:_ -> isLower c

unexpected :: String -> KindError
unexpected message = 
   internalError "KindInferencing.ag" "unexpected" ("unexpected kind error: "++message)

(<==>) :: Kind -> Kind -> ((Kind, Kind) -> KindError) -> KindConstraint
(k1 <==> k2) info = (k1 .==. k2) (info (k1, k2))
}