{-# 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 /parametric polymorphism/.
module Language.LOL.Typing.Solver.Polytype where

import Control.Monad (Monad(..), mapM, forM_, foldM, sequence)
import qualified Control.Monad.Classes as MC
import Data.Bool
import Data.Eq (Eq(..))
import Data.Function (($), flip)
import Data.Functor ((<$>))
import qualified Data.List as List
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..), (<>))
import Data.Text.Buildable (Buildable(..))
import Data.Tuple (snd)
import Prelude (Num(..), error)
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 qualified Language.LOL.Typing.Lib.Control.Monad.Classes.Instance as MC
import Language.LOL.Typing.Lib.Data.Default (Default(..))
import Language.LOL.Typing.Lib.Data.Empty (Empty(..))
import qualified Language.LOL.Typing.Lib.Data.Text.Buildable as Build

-- * Type 'State_Polytype'

data State_Polytype info
 =   State_Polytype
 {   state_polytype_freshvar :: Freshvar
     -- ^ An incrementing counter
     -- for fresh 'Monovar's, fresh 'Rigvar's, or 'Polyvar's.
 ,   state_polytype_substitution :: Polysub
     -- ^ Known 'Polytype's.
 ,   state_polytype_rigids :: [Infoed info ([Rigvar], [Rigtype])]
     -- ^ The 'Rigvar's introduced as 'Freshvar's by 'polytype_rigvarify'
     -- with the 'Rigtype's from the /typing context/
     -- (used by 'polytype_rigids_check').
 , state_polytype_quantifiers :: Map Monovar [Quantifier]
     -- ^ Map every 'Monovar' unified with an instance of a 'Polytype'
     -- to the 'Freshvar's allocated to instantiate
     -- the 'quantifiers' of this 'Polytype'.
 } deriving (Show)

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

instance Buildable info => Buildable (State_Polytype info) where
	build State_Polytype
	 { state_polytype_freshvar     = freshvar
	 , state_polytype_substitution = polytys
	 , state_polytype_rigids       = rigids
	 , state_polytype_quantifiers    = instances
	 } =
		Build.unlines $
		 [ "freshvar: " <> build freshvar
		 , "rigids: " ] <>
			(("  " <>) <$>
			List.foldr (\(Infoed info (rigvas, monotys)) ->
				((Build.tuple
				 [ Build.list rigvas
				 , Build.list monotys
				 ] <> " {- " <> build info <> " -}") :)
			 ) [] rigids) <>
		 [ "polytypes: " ] <>
			Map.foldrWithKey (\v ty ->
				(("  p" <> build v <> " == " <>
				build (def::Quantification_Build_Options, ty)) :)
			 ) [] polytys <>
		 [ "instances: " ] <>
			Map.foldrWithKey (\v vs ->
				(("  " <> build (Monotype_Var v) <> " -> " <>
				Build.list (Monotype_Var <$> vs)) :)
			 ) [] instances
instance Buildable info => State (State_Polytype info) where
	state_name _ = "State_Polytype"
	state_show = Build.text
instance Empty (State_Polytype m) where
	empty =
		State_Polytype
		 { state_polytype_freshvar     = 0
		 , state_polytype_substitution = mempty
		 , state_polytype_rigids       = mempty
		 , state_polytype_quantifiers    = mempty
		 }

-- * Class 'Solver_Polytype'

class
 ( Solver_Constraint m
 , MC.MonadState (State_Polytype (Info m)) m
 , Infoable Info_Polytype (Info m)
 , Solver_Logable Log_Polytype m
 ) => Solver_Polytype m where
	error_polytype :: Error_Polytype -> Error m
	
	polytype_freshvar :: m Freshvar
	polytype_freshvar =
		MC.gets $ \(s::State_Polytype (Info m)) ->
			state_polytype_freshvar s
	polytype_freshvar_set :: Freshvar -> m ()
	polytype_freshvar_set state_polytype_freshvar =
		MC.modify $ \(s::State_Polytype (Info m)) ->
			s{ state_polytype_freshvar }
	
	polytype_substitution :: m Polysub
	polytype_substitution =
		MC.gets $ \(s::State_Polytype (Info m)) ->
			state_polytype_substitution s
	polytype_insert :: Polyvar -> Polytype -> m ()
	polytype_insert var polyty = do
		log $ Log_Polytype_Insert var polyty
		MC.modify $ \(s::State_Polytype (Info m)) ->
			s{ state_polytype_substitution =
				Map.insert var polyty $
				state_polytype_substitution s
			 }
	polytype_lookup
	 :: Polytyref -> m Polytype
	polytype_lookup polyref =
		case polyref of
		 Polytyref polyty -> return polyty
		 Polytyref_Var v ->
			Map.findWithDefault
			 (error "polytype_lookup: Polytyref_Var not found in polytype_substitution")
			 v <$> polytype_substitution
	
	-- | 'Rigvar' variables
	--
	-- NOTE: used by 'polytype_rigids_check'.
	polytype_rigids :: m [Infoed (Info m) ([Rigvar], [Rigtype])]
	polytype_rigids =
		MC.gets $ \(s::State_Polytype (Info m)) ->
			state_polytype_rigids s
	polytype_rigids_set :: [Infoed (Info m) ([Rigvar], [Rigtype])] -> m ()
	polytype_rigids_set state_polytype_rigids =
		MC.modify $ \(s::State_Polytype (Info m)) ->
			s{ state_polytype_rigids }
	polytype_rigids_push
	 :: Infoed (Info m) ([Rigvar], [Rigtype]) -> m ()
	polytype_rigids_push rig{-@(Infoed info (vs, ms))-} = do
		--constraint_log "polytype_rigids_push"
		-- (Infoed info (Build.tuple [Build.list vs, Build.list ms]))
		rigs <- polytype_rigids
		polytype_rigids_set (rig:rigs)
	
	-- | Return the 'quantified' of the given 'Forall'
	-- with the 'quantifiers' 'substitute'd with 'Freshvar's.
	-- Also, if the given 'Monotype' is a 'Monotype_Var',
	-- and at least one 'Freshvar' has been allocated,
	-- record them in the 'State_Polytype'.
	polytype_instantiate
	 :: Substitutable a => Monotype -> Forall a -> m a
	polytype_instantiate monoty quant = do
		freshvar <- polytype_freshvar
		let Unquantification freshvar' a = unquantify_forall freshvar quant
		polytype_freshvar_set freshvar'
		case monoty of
		 Monotype_Var monovar ->
			case [freshvar..freshvar'-1] of
			 [] -> return ()
			 fs ->
				MC.modify $ \(s::State_Polytype (Info m)) ->
					s{ state_polytype_quantifiers =
						Map.insert monovar fs $
						state_polytype_quantifiers s
					 }
		 _ -> return ()
		return a
	polytype_quantifiers :: m (Map Monovar [Quantifier])
	polytype_quantifiers =
		MC.gets $ \(s::State_Polytype (Info m)) ->
			state_polytype_quantifiers s
	
	-- | Like 'rigvarify', but 'unquantify'
	-- by using fresh 'Monovar's ('Freshvar') instead of fresh 'Const's ('Rigvar's).
	--
	-- NOTE: useful for the 'quantifiers' of a 'Polytype',
	-- maintaining a list of 'Rigvar's in 'state_polytype_rigids'.
	polytype_rigvarify
	 :: Substitutable a
	 => Info m -> [Rigtype] -> Forall a -> m a
	polytype_rigvarify info rigtys forall_a = do
		freshvar <- polytype_freshvar
		let Unquantification freshvar' a =
			unquantify_forall freshvar forall_a
		polytype_freshvar_set freshvar'
		polytype_rigids_push $
			Infoed info ([freshvar .. freshvar'-1], rigtys)
		return a

polytype_rigids_check
 :: forall m.
 ( Solver_Constraint m
 , Solver_Monotype m
 , Solver_Polytype m
 ) => m ()
polytype_rigids_check =
	polytype_rigids >>=
		rigvars_substitution >>=
		check_rigvars_are_mapped_to_vars >>=
		check_rigvars_are_mapped_to_different_vars >>=
		check_monotys_have_no_rigvar >>=
		rigvars_substituted >>=
		polytype_rigids_set
		 -- NOTE: restore the 'Rigvar's
		 -- which are consistent with the current 'monotype_substitution'.
	where
		-- | Return the 'Monotype's toward which 'monotype_substitution'
		-- maps the given 'Rigvar's.
		rigvars_substitution
		 ::   [Infoed (Info m) ([Rigvar]            , [Rigtype])]
		 -> m [Infoed (Info m) ([(Rigvar, Monotype)], [Rigtype])]
		rigvars_substitution =
			mapM $ \(Infoed info (rigvas, rigtys)) -> do
				subtys <- monotype_lookup `mapM` rigvas
				return $ Infoed info
				 ( List.zip rigvas subtys
				 , rigtys )
		-- | Return the 'Monovar's of the given 'Monotype's which are 'Monotype_Var's.
		check_rigvars_are_mapped_to_vars
		 ::   [Infoed (Info m) ([(Rigvar, Monotype)], [Rigtype])]
		 -> m [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		check_rigvars_are_mapped_to_vars tochecks = do
			let (goods, bads) =
				List.foldr
				 (\i@(Infoed info (rigvars_monotys, rigtys)) (gs, bs) ->
					let (rigvars_subvars::Maybe [(Rigvar, Monovar)]) =
						sequence $ (\(rig, ty) ->
							case ty of
							 Monotype_Var v -> Just (rig, v)
							 _ -> Nothing
						 ) <$> rigvars_monotys in
					case rigvars_subvars of
					 Just x  -> (Infoed info (x, rigtys):gs, bs)
					 Nothing -> (gs, i:bs)
				 )
				 ([], [])
				 tochecks
			forM_ bads $ \(Infoed info (rigvars_monotys, _monotys)) ->
				constraint_error_insert
				 (error_polytype $ Error_Polytype_Rigid_type_mismatch rigvars_monotys)
				 info
			return goods
		-- | Return the 'Rigvar's which are injectively mapped to the given 'Monovar's.
		check_rigvars_are_mapped_to_different_vars
		 ::   [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		 -> m [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		check_rigvars_are_mapped_to_different_vars tochecks = do
			let (subvars_rigvars::Map Monovar [Rigvar]) =
				Map.fromListWith (flip (<>))
				 [ (subvar, [rigvar])
				 | Infoed _ (rigvars_subvars, _) <- tochecks
				 , (rigvar, subvar) <- rigvars_subvars
				 ]
			let subvars_colliding_rigvars =
				Map.filter
				 (\l -> case l of { [] -> False; [_] -> False; _ -> True })
				 subvars_rigvars
			let (goods, bads) =
				List.partition
				 (\(Infoed _ (rigvars_subvars, _)) ->
					List.null $ (subvars_colliding `List.intersect`) $
					snd <$> rigvars_subvars)
				 tochecks
				where subvars_colliding = Map.keys subvars_colliding_rigvars
			forM_ bads $ \(Infoed info (rigvars_subvars, _)) ->
				constraint_error_insert
				 (error_polytype $ Error_Polytype_Rigid_injectivity_lost $
					Map.fromList
					 [ (subvar, rigvas)
					 | (_rigvar, subvar) <- rigvars_subvars
					 , Just rigvas <- [Map.lookup subvar subvars_colliding_rigvars]
					 ])
				 info
			return goods
		check_monotys_have_no_rigvar
		 ::   [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		 -> m [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		check_monotys_have_no_rigvar =
			foldM (\goods this@(Infoed info (rigvars_subvars, rigtys)) -> do
				rigtys' <- monotype_substitute rigtys
				case subvars rigtys' `List.intersect` (snd <$> rigvars_subvars) of
				 []  -> return (this:goods)
				 rigvars_escaping -> do
					constraint_error_insert
					 (error_polytype $ Error_Polytype_Rigid_escaping rigvars_escaping)
					 info
					return goods
			 ) []
		rigvars_substituted
		 ::   [Infoed (Info m) ([(Rigvar, Monovar)], [Rigtype])]
		 -> m [Infoed (Info m) ([Rigvar], [Rigtype])]
		rigvars_substituted checkeds =
			return
			 [ Infoed info (snd <$> rigvars_subvars, rigtys)
			 | Infoed info (        rigvars_subvars, rigtys) <- checkeds
			 ]

-- * Class 'Info_Polytype'

data Info_Polytype
 =   Info_Polytype_Instantiated Polytype
 |   Info_Polytype_Rigidified [Rigtype] Polytype
 deriving (Eq, Show)

instance Buildable Info_Polytype where
	build x =
		case x of
		 Info_Polytype_Instantiated p ->
			"Info_Polytype_Instantiated " <> build p
		 Info_Polytype_Rigidified rs p ->
			"Info_Polytype_Rigidified " <> Build.list rs <> " " <> build p

-- * Type 'Error_Polytype'

data Error_Polytype
 =   Error_Polytype_Rigid_type_mismatch [(Rigvar, Monotype)]
     -- ^ A 'Rigvar' is mapped by the current 'monotype_substitution'
     -- to a 'Monotype' which is not a 'Monotype_Var'.
 |   Error_Polytype_Rigid_injectivity_lost (Map Monovar [Rigvar])
     -- ^ 'Map' of 'Monovar's reached by at least two distincts 'Rigvar's.
 |   Error_Polytype_Rigid_escaping [Rigvar]
     -- ^ Some 'Rigvar's escape via a 'Rigtype'.
 deriving (Eq, Show)

-- ** Type 'Log_Polytype'

data Log_Polytype
 =   Log_Polytype_Insert Polyvar Polytype
 deriving (Show)
instance Buildable Log_Polytype where
	build x =
		case x of
		 Log_Polytype_Insert var polyty ->
			"polytype_insert : " <>
			build (Polytyref_Var var) <> " == " <> build polyty