{-# LANGUAGE GeneralizedNewtypeDeriving, LambdaCase #-}
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
module Type 
  ( module Type
  , Id(..)
  ) where

import Control.Monad.RWS
import Control.Monad.Reader
import Control.Monad.Except

import Data.Map (Map)
import Data.Set (Set)
import Data.Text (Text)

import qualified Data.Set as S
import qualified Data.Map as M
import qualified Data.Text as T

import Hm.Abs (Id(..))

import qualified Hm.Abs as H (TypeSig'(..), TypeSig)

import Misc

import Prelude hiding (map)

data PolyT
  = Forall (Set Id) MonoT -- ^ ∀ σ₁ σ₂ … σₙ. τ
  | Mono MonoT        -- ^ τ
  deriving Show

infixr `TArr`
data MonoT
  = TArr MonoT MonoT -- ^ function
  | TVar Id          -- ^ variable
  | TCon Id          -- ^ constant
  deriving (Eq, Show)

type Pos = Maybe (Int, Int)

type TL = TL' Pos
data TL' a
  = TypeDef a Id [Id]  TypeEnv -- ^ name, parameters, constructors and recursor
  | VarDef  a Id PolyT Exp     -- ^ name, declared type, expression
  deriving Show

type Exp = Exp' Pos
data Exp' a
  = Let a [(Id,Exp)] Exp
  | Abs a [Id] Exp
  | App a Exp [Exp]
  | Var a Id
  deriving Show

class Positioned p where
  pos :: p -> Pos

instance Positioned TL where
  pos = \case
    TypeDef p _ _ _ -> p
    VarDef  p _ _ _ -> p

instance Positioned Exp where
  pos = \case
    Let   p _ _ -> p
    Abs   p _ _ -> p
    App   p _ _ -> p
    Var   p _   -> p

instance Positioned Pos where
  pos = id

instance Positioned H.TypeSig where
  pos = \case
    H.TypeFun p _ _ -> p
    H.TypeApp p _ _ -> p
    H.TypeVar p _   -> p

class Substitutable a where
  apply :: Subst -> a -> a -- ^ apply a substitution
  free   :: a -> Set Id    -- ^ free type variables

instance Substitutable MonoT where
  apply s = \case
    TCon i -> TCon i
    TVar i -> lookupDefault (TVar i) i s
    (t1 `TArr` t2) -> apply s t1 `TArr` apply s t2

  free = \case
    TCon{} -> S.empty
    TVar i -> S.singleton i
    (t1 `TArr` t2) -> free t1 <> free t2

instance Substitutable PolyT where
  apply s = \case
    Forall as t -> Forall as (apply (foldr M.delete s as) t)
    Mono t      -> Mono (apply s t)

  free = \case
    Forall as t -> free t \\ as
    Mono t -> free t

instance Substitutable TypeEnv where
  apply s = map (apply s)
  free    = free . M.elems

instance Substitutable a => Substitutable [a] where
  apply = map . apply
  free  = foldMap free

instance Substitutable Constraint where
  apply s (a, b, d) = (apply s a, apply s b, d)
  free    (a, b, _) = free a <> free b

data TypeError
  = Oop -- ^ compiler error (oops)
  | UnificationFailure MonoT MonoT
  | UnificationRight MonoT MonoT
  | InfiniteType Id MonoT
  | UnboundVariable Pos Id
  | Unimplemented Text
  | InvalidTypeDecl
  | InvalidConstructor
  | ArityMismatch
  | AlreadyDefined Pos Id
  | MutuallyRecursive (Set Id)
  | PositivityCheck Id
  deriving Show

type TypeEnv = Map Id PolyT

type Subst = Map Id MonoT

emptySubst :: Subst
emptySubst = M.empty

-- This substution, and that one
(<&>) :: Subst -> Subst -> Subst
(<&>) s1 s2 = map (apply s1) s2 <> s1

data CT
  = Unify
  | UnifyRight
  deriving (Eq, Show)

type Constraint = (MonoT, MonoT, CT)

type CheckState = [Id] 

initialState :: [Id]
initialState = [1..] >>= map (Id . T.pack) . flip replicateM ['a'..'z']

newtype Infer a = Infer { getInfer :: RWST TypeEnv [Constraint] CheckState (Except TypeError) a }
  deriving ( Functor, Applicative, Monad
           , MonadError TypeError
           , MonadState CheckState
           , MonadReader TypeEnv
           , MonadWriter [Constraint]
           )

instance MonadFail Infer where
  fail _ = throwError Oop

type Unifier = (Subst, [Constraint])

emptyUnifier :: Unifier
emptyUnifier = (emptySubst, [])

newtype Solve a = Solve { getSolve :: ReaderT Unifier (Except TypeError) a}
  deriving ( Functor, Applicative, Monad
           , MonadError TypeError
           , MonadReader Unifier
           )

instance MonadFail Solve where
  fail _ = throwError Oop

type RefNode  = (Id, Set Id)
type RefGraph = Map Id (Set Id)