{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}
-- | Type inference for /type class polymorphism/
-- (aka. /parametric overloading/).
module Language.LOL.Typing.Solver.Class where

import Control.Monad (Monad(..), forM, forM_, mapM)
import qualified Control.Monad.Classes as MC
import Data.Bool
import Data.Eq (Eq(..))
import qualified Data.Foldable as Foldable
import Data.Function (($), (.))
import Data.Functor ((<$>))
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..), isJust)
import Data.Monoid (Monoid(..), (<>))
import Data.Text (Text)
import Data.Text.Buildable (Buildable(..))
import Data.Tuple (fst)
import Text.Show (Show(..))

import Language.LOL.Typing.Type
import Language.LOL.Typing.Solver.Common
import Language.LOL.Typing.Solver.Constraint
import Language.LOL.Typing.Solver.Monotype
import Language.LOL.Typing.Solver.Polytype
import qualified Language.LOL.Typing.Lib.Control.Monad.Classes.Instance as MC
import Language.LOL.Typing.Lib.Data.Empty (Empty(..))
import qualified Language.LOL.Typing.Lib.Data.Text.Buildable as Build

-- * Type 'State_Class'

data State_Class info
 =   State_Class
 {   state_class_directives :: [Class_Directive info]
     -- ^ 'Class_Directive's for 'state_class_qualifiers'.
 ,   state_class_env :: Class_Env
     -- ^ Known 'Class'es and 'Class_Instance's.
 ,   state_class_qualifiers :: State_Class_Qualifiers info
     -- ^ 'Class' assertions.
 } deriving (Show)

-- | Make 'State_Class' collectable as a 'State' instance
-- out of a 'Monad' stack.
type instance MC.Class State (State_Class info) = 'True

instance Empty (State_Class info) where
	empty = State_Class
	 { state_class_directives = mempty
	 , state_class_env        = mempty
	 , state_class_qualifiers = empty
	 }
instance Buildable info => State (State_Class info) where
	state_name _ = "State_Class"
	state_show = Build.text
instance Buildable info => Buildable (State_Class info) where
	build State_Class
	 { state_class_directives
	 , state_class_env
	 , state_class_qualifiers
	 } =
		mconcat $ List.intersperse "\n"
		 [ "class directives: "  <> Build.list state_class_directives
		 , "class environment: " <> Build.list (Map.keys state_class_env)
		 , "class qualifiers: "
		 , Build.indent "  " state_class_qualifiers
		 ]

-- ** Type 'State_Class_Qualifiers'

data State_Class_Qualifiers info
 =   State_Class_Qualifiers
 {   state_class_qualifiers_assumed     :: [Infoed info Class_Qualifier]
 ,   state_class_qualifiers_generalized :: [Infoed info Class_Qualifier]
 ,   state_class_qualifiers_toprove     :: [Infoed info Class_Qualifier]
     -- ^ Once the constraints have been solved, this list should be empty.
 } deriving (Show)
 
instance Empty (State_Class_Qualifiers info) where
	empty =
		State_Class_Qualifiers
		 { state_class_qualifiers_assumed     = []
		 , state_class_qualifiers_generalized = []
		 , state_class_qualifiers_toprove     = []
		 }
instance Buildable info => Buildable (State_Class_Qualifiers info) where
	build State_Class_Qualifiers
	 { state_class_qualifiers_assumed
	 , state_class_qualifiers_generalized
	 , state_class_qualifiers_toprove
	 } =
		Build.unlines
		 [ "assumed: "     <> Build.list state_class_qualifiers_assumed
		 , "generalized: " <> Build.list state_class_qualifiers_generalized
		 , "toprove: "     <> Build.list state_class_qualifiers_toprove
		 ]
instance Substitutable (State_Class_Qualifiers info) where
	subvars (State_Class_Qualifiers qs gs as) = subvars (infoed <$> (qs <> gs <> as))
	sub `substitute` State_Class_Qualifiers
	 { state_class_qualifiers_toprove     = ps
	 , state_class_qualifiers_generalized = gs
	 , state_class_qualifiers_assumed     = as
	 } =
		State_Class_Qualifiers
		 { state_class_qualifiers_toprove     = ((sub `substitute`) <$>) <$> ps
		 , state_class_qualifiers_generalized = ((sub `substitute`) <$>) <$> gs
		 , state_class_qualifiers_assumed     = ((sub `substitute`) <$>) <$> as
		 }

-- * Class 'Solver_Class'

class
 ( Solver_Constraint m
 , Solver_Monotype m
 , Solver_Polytype m
 , MC.MonadState (State_Class (Info m)) m
 , Infoable (Info_Class (Info m)) (Info m)
 , Solver_Logable Log_Class m
 ) => Solver_Class m where
	error_class :: Error_Class -> Error m
	
	class_qualifiers :: m (State_Class_Qualifiers (Info m))
	class_qualifiers =
		MC.gets $ \(s::State_Class (Info m)) ->
			state_class_qualifiers s
	class_qualifiers_modify
	 :: (State_Class_Qualifiers (Info m) -> State_Class_Qualifiers (Info m))
	 -> m ()
	class_qualifiers_modify f =
		MC.modify $ \(s::State_Class (Info m)) ->
			s{ state_class_qualifiers =
				f (state_class_qualifiers s) }
	class_qualifier_toprove
	 :: Info m -> Class_Qualifier -> m ()
	class_qualifier_toprove info q = do
		log $ Log_Class_Qualifier_ToProve q
		class_qualifiers_modify $ \quals ->
			quals{ state_class_qualifiers_toprove =
				Infoed info q : state_class_qualifiers_toprove quals }
	class_qualifier_assume
	 :: Info m -> Class_Qualifier -> m ()
	class_qualifier_assume info q = do
		log $ Log_Class_Qualifier_Assume q
		class_qualifiers_modify $ \quals ->
			quals{ state_class_qualifiers_assumed =
				Infoed info q : state_class_qualifiers_assumed quals }
	class_qualifiers_map
	 :: (Class_Qualifier -> m Class_Qualifier) -> m ()
	class_qualifiers_map f = do
		let g = mapM $ \(Infoed info q) ->
			f q >>= \new -> return (Infoed info new)
		state_class_qualifiers_toprove <-
			(>>= g) $
			MC.gets $ \(s::(State_Class (Info m))) ->
				state_class_qualifiers_toprove $
				state_class_qualifiers s
		state_class_qualifiers_generalized <-
			(>>= g) $
			MC.gets $ \(s::(State_Class (Info m))) ->
				state_class_qualifiers_generalized $
				state_class_qualifiers s
		state_class_qualifiers_assumed <-
			(>>= g) $
			MC.gets $ \(s::(State_Class (Info m))) ->
				state_class_qualifiers_assumed $
				state_class_qualifiers s
		class_qualifiers_modify $ \quals ->
			quals
			 { state_class_qualifiers_toprove
			 , state_class_qualifiers_generalized
			 , state_class_qualifiers_assumed
			 }
	
	class_qualifiers_reduced :: m [Class_Qualifier]
	class_qualifiers_reduced = do
		synotys <- synotype_substitution
		clenv <- class_env
		quals <- class_qualifiers
		return $ fst $
			class_context_reduction
			 synotys clenv $ infoed <$> (
				state_class_qualifiers_toprove quals <>
				state_class_qualifiers_generalized quals <>
				state_class_qualifiers_assumed quals
			 )
	
	-- | Generalize a 'Monotype', preserving given 'Rigtype's,
	-- and such that the resulting 'Polytype'
	-- has all the 'Class_Qualifier's from 'state_class_qualifiers_toprove',
	-- when they apply at least on one of the 'quantifiers' of this 'Polytype'.
	class_polytype_forall
	 :: [Rigtype] -> Monotype -> m Polytype
	class_polytype_forall rigtys monoty = do
		State_Class_Qualifiers
		 { state_class_qualifiers_toprove     = quals_toprove
		 , state_class_qualifiers_generalized = quals_generalized
		 } <- class_qualifiers
		let rigvas = subvars rigtys
		let polyvars = subvars monoty List.\\ rigvas
		let has_polyvars = Foldable.any (`List.elem` polyvars) . subvars . infoed
		let (  quals_toprove_poly
		     , quals_toprove_mono
		     ) = List.partition has_polyvars quals_toprove
		let quals_generalized_poly = List.filter has_polyvars quals_generalized
		class_qualifiers_modify $ \quals ->
			quals
			 { state_class_qualifiers_toprove     = quals_toprove_mono
			 , state_class_qualifiers_generalized = quals_toprove_poly <> state_class_qualifiers_generalized quals
			 }
		return $
			forall_but rigvas $
			Qualification
			 (infoed <$> (quals_toprove_poly <> quals_generalized_poly))
			 monoty
	
	class_improve
	 :: Bool -> m [(Info m, Monotype, Monotype)]
	class_improve normal =
		if normal
		then class_improve_normal
		else class_improve_final
	class_improve_normal
	 :: m [(Info m, Monotype, Monotype)]
	class_improve_normal =
		return []
	class_improve_final
	 :: m [(Info m, Monotype, Monotype)]
	class_improve_final =
		return []
	
	class_simplify :: m ()
	class_simplify = do
		State_Class_Qualifiers
		 { state_class_qualifiers_toprove
		 , state_class_qualifiers_assumed
		 } <- class_qualifiers
		synotys <- synotype_substitution
		clenv <- class_env
		cldirs <- MC.gets state_class_directives
		clquals <- simplify synotys clenv cldirs
			 state_class_qualifiers_toprove
		class_qualifiers_modify $ \quals ->
			quals{state_class_qualifiers_toprove =
				List.filter (not . class_entails synotys clenv
				 (infoed <$> state_class_qualifiers_assumed) . infoed) clquals }
		where
		simplify ::
		 ( Solver_Constraint m
		 , Solver_Polytype m
		 ) => Synotype_Substitution
		 -> Class_Env
		 -> [Class_Directive (Info m)]
		 ->   [Infoed (Info m) Class_Qualifier]
		 -> m [Infoed (Info m) Class_Qualifier]
		simplify synotys clenv cldirs clquals = do
			hnf' <- go_insts clquals
			are_disjoints (go_super_class [] hnf')
			where
			go_insts
			 ::   [Infoed (Info m) Class_Qualifier]
			 -> m [Infoed (Info m) Class_Qualifier]
			go_insts ts =
				(List.concat <$>) <$>
				forM ts $
				 \q@(Infoed info clqual@(Class_Qualifier clname _)) ->
					if is_class_qualifier_normalized clqual
					then return [q]
					else case class_env_instance_context synotys clenv clqual of
					 Just inst_context ->
						go_insts $
						 Infoed (info_insert
							 (Info_Class_Qualifier_Parent clqual::Info_Class (Info m))
							 info) <$>
						 inst_context
					 Nothing -> do
						constraint_error_insert
						 (error_class Error_Class_Qualifier_Unresolved) $
						 (case cldirs_never of
							 clql:_ -> info_insert $ Info_Class_Directive_Never clql
							 [] -> case infos_cldirs_close of
								 [i] -> info_insert (Info_Class_Directive_Close (Infoed i clname)::Info_Class (Info m))
								 _   -> info_insert (Info_Class_Qualifier_Unresolved clqual::Info_Class (Info m)))
						 info
						return []
						where
						cldirs_never =
							[ Infoed i clql
							| Class_Directive_Never clql i <- cldirs
							, isJust $ class_qualifier_unification synotys clqual clql
							]
						infos_cldirs_close =
							[ i
							| Class_Directive_Close n i <- cldirs
							, n == clname ]
			
			go_super_class
			 :: [Infoed (Info m) Class_Qualifier]
			 -> [Infoed (Info m) Class_Qualifier]
			 -> [Infoed (Info m) Class_Qualifier]
			go_super_class rs [] = rs
			go_super_class rs (x:xs)
			 | class_entails_super_class clenv
				 (infoed <$> (rs <> xs)) (infoed x) = go_super_class rs xs
			 | otherwise = go_super_class (x:rs) xs
			
			are_disjoints
			 ::   [Infoed (Info m) Class_Qualifier]
			 -> m [Infoed (Info m) Class_Qualifier]
			are_disjoints [] = return []
			are_disjoints (t@(Infoed info (Class_Qualifier className ty)):ts) = do
				let f t'@(Infoed info' (Class_Qualifier className' ty')) =
					case
					 [ i
					 | ty == ty'
					 , Class_Directive_Disjoint ss i <- cldirs
					 , className  `List.elem` ss
					 , className' `List.elem` ss
					 ] of
					 [] -> return ([t'], True)
					 info_directive : _ -> do
						constraint_error_insert
						 (error_class Error_Class_Qualifier_Disjoint) $
							info_insert
							 (Info_Class_Directive_Disjoint
								 (Infoed info className)
								 (Infoed info' className'))
							 info_directive
						return ([], False)
				result <- mapM f ts
				let (list, bs) = List.unzip result
				rest <- are_disjoints (List.concat list)
				return $ if Foldable.and bs then t : rest else rest
	
	class_ambiguous :: m ()
	class_ambiguous = do
		State_Class_Qualifiers
		 { state_class_qualifiers_toprove
		 } <- class_qualifiers
		rigvars_rigtypes <- polytype_rigids
		forM_ state_class_qualifiers_toprove $ \q_toprove ->
			case q_toprove of
			 Infoed _ Class_Qualifier{class_qualifier_type=Monotype_Var v} ->
				case
				 [ info
				 | Infoed info (rigvas, _) <- rigvars_rigtypes
				 , v `List.elem` rigvas
				 ] of
				 info:_ -> err_missing q_toprove info
				 _      -> err_ambiguous q_toprove
			 _ -> err_ambiguous q_toprove
		where
		err_ambiguous (Infoed info p) =
			constraint_error_insert
			 (error_class Error_Class_Qualifier_Ambiguous) $
			info_insert
			 (Info_Class_Qualifier_Ambiguous p::Info_Class (Info m)) info
		err_missing q_toprove info =
			constraint_error_insert
			 (error_class Error_Class_Qualifier_Missing) $
			info_insert
			 (Info_Class_Qualifier_Arising_from q_toprove) info
	
	class_env :: m Class_Env
	class_env =
		MC.gets $ \(s::State_Class (Info m)) ->
			state_class_env s
	class_env_set :: Class_Env -> m ()
	class_env_set state_class_env =
		MC.modify $ \(s::State_Class (Info m)) ->
			s{ state_class_env }

-- | When all 'state_constraint_constraints' have been handled,
-- takes all the remaining 'state_class_qualifiers_toprove',
-- and report them as ambiguities.
class_ambiguities ::
 ( Solver_Constraint m
 , Solver_Monotype m
 , Solver_Polytype m
 , Solver_Class m
 ) => m ()
class_ambiguities = do
	class_reduction
	class_improve_fix False
	class_ambiguous

-- | Perform context reduction on the 'state_class_qualifiers_toprove',
-- and removes the 'Class_Qualifier's
-- which are entailed by 'state_class_qualifiers_assumed'.
class_reduction ::
 ( Solver_Constraint m
 , Solver_Monotype m
 , Solver_Polytype m
 , Solver_Class m
 ) => m ()
class_reduction = do
	monotype_substitution_consistentify
	class_qualifiers_map monotype_substitute
	class_improve_fix True
	class_simplify

class_improve_fix ::
 ( Solver_Constraint m
 , Solver_Monotype m
 , Solver_Polytype m
 , Solver_Class m
 ) => Bool -> m ()
class_improve_fix normal = do
	improvements <- class_improve normal
	case improvements of
	 [] -> return ()
	 _  -> do
		forM_ improvements
		 (\(info, t1, t2) -> monotype_unify info t1 t2)
		monotype_substitution_consistentify
		class_improve_fix normal

-- * Class 'Info_Class'

data Info_Class info
 =   Info_Class_Qualifier_Ambiguous Class_Qualifier
 |   Info_Class_Qualifier_Arising_from (Infoed info Class_Qualifier)
 |   Info_Class_Qualifier_Parent Class_Qualifier
 |   Info_Class_Qualifier_Unresolved Class_Qualifier
 |   Info_Class_Directive_Close (Infoed info Text)
 |   Info_Class_Directive_Disjoint (Infoed info Text) (Infoed info Text)
 |   Info_Class_Directive_Never (Infoed info Class_Qualifier)
 deriving (Eq, Show)

instance Buildable info => Buildable (Info_Class info) where
	build x =
		case x of
		 Info_Class_Qualifier_Ambiguous q ->
			"Info_Class_Qualifier_Ambiguous " <> build q
		 Info_Class_Qualifier_Arising_from q ->
			"Info_Class_Qualifier_Arising_from " <> build q
		 Info_Class_Qualifier_Parent q ->
			"Info_Class_Qualifier_Parent " <> build q
		 Info_Class_Qualifier_Unresolved q ->
			"Info_Class_Qualifier_Unresolved " <> build q
		 Info_Class_Directive_Close n ->
			"Info_Class_Directive_Close " <> build n
		 Info_Class_Directive_Disjoint n1 n2 ->
			"Info_Class_Directive_Disjoint " <> build n1 <> " " <> build n2
		 Info_Class_Directive_Never q ->
			"Info_Class_Directive_Never " <> build q

-- * Type 'Error_Class'

data Error_Class
 =   Error_Class_Qualifier_Ambiguous
     -- ^ Example: @forall a. Eq a => Int -> Int@
 |   Error_Class_Qualifier_Disjoint
 |   Error_Class_Qualifier_Missing
 |   Error_Class_Qualifier_Unresolved
     -- ^ A 'Monotype' is missing a required 'Class_Qualifier'.
 deriving (Eq, Show)

-- ** Type 'Log_Class'

data Log_Class
 =   Log_Class_Qualifier_Assume  Class_Qualifier
 |   Log_Class_Qualifier_ToProve Class_Qualifier
 deriving (Show)
instance Buildable Log_Class where
	build x =
		case x of
		 Log_Class_Qualifier_Assume  q -> "class_qualifier_assume  : " <> build q
		 Log_Class_Qualifier_ToProve q -> "class_qualifier_toprove : " <> build q