{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE ViewPatterns               #-}

------------------------------------------------------------------------------
-- The MIT License (MIT)
--
-- Copyright (c) 2017 Luka Horvat, 2019 Sandy Maguire
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to
-- deal in the Software without restriction, including without limitation the
-- rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
-- sell copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
--
-- The above copyright notice and this permission notice shall be included in
-- all copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
-- FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
-- IN THE SOFTWARE.
--
------------------------------------------------------------------------------
--
-- This module was originally based on 'Control.Effects.Plugin' from the
-- 'simple-effects' package, by Luka Horvat.
--
-- https://gitlab.com/LukaHorvat/simple-effects/commit/966ce80b8b5777a4bd8f87ffd443f5fa80cc8845#f51c1641c95dfaa4827f641013f8017e8cd02aab

module Polysemy.Plugin.Fundep (fundepPlugin) where

import           Control.Monad
import           Control.Monad.Trans.Class (lift)
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Coerce
import           Data.Function (on)
import           Data.IORef
import qualified Data.Map as M
import           Data.Maybe
import           Data.Set (Set)
import qualified Data.Set as S
import           Data.Traversable (for)
import           Polysemy.Plugin.Fundep.Stuff
import           Polysemy.Plugin.Fundep.Unification
import           Polysemy.Plugin.Fundep.Utils

#if __GLASGOW_HASKELL__ >= 906
#define SUBST Subst
#define GET_TYVAR getTyVar_maybe
import           GHC.Core.TyCo.Subst (SUBST)
import           GHC.Core.TyCo.Compare (eqType, nonDetCmpType)
#else
#define SUBST TCvSubst
#define GET_TYVAR tcGetTyVar_maybe
#endif

#if __GLASGOW_HASKELL__ >= 900
import           GHC.Builtin.Types.Prim (alphaTys)
import           GHC.Plugins (idType, tyConClass_maybe, ppr, Outputable, sep, text, (<+>), parens, emptyUFM)
import           GHC.Tc.Types.Evidence
import           GHC.Tc.Plugin (TcPluginM, tcPluginIO)
import           GHC.Tc.Types
import           GHC.Tc.Types.Constraint
import           GHC.Tc.Utils.Env (tcGetInstEnvs)
import           GHC.Tc.Utils.TcType (tcSplitPhiTy, tcSplitTyConApp, GET_TYVAR, tcSplitAppTy_maybe)
import           GHC.Tc.Solver.Monad hiding (tcLookupClass)
import           GHC.Core.Class (classTyCon)
import           GHC.Core.InstEnv (lookupInstEnv, is_dfun)
import           GHC.Core.Type
import           GHC.Utils.Monad (allM, anyM)

#else
#if __GLASGOW_HASKELL__ >= 810
import           Constraint
#endif

import           Class (classTyCon)
import           GhcPlugins (idType, tyConClass_maybe, ppr, Outputable, sep, text, (<+>), parens)
import           Inst (tcGetInstEnvs)
import           InstEnv (lookupInstEnv, is_dfun)
import           MonadUtils (allM, anyM)
import           TcEvidence
import           TcPluginM (tcPluginIO)
import           TcRnTypes
import           TcType (tcSplitPhiTy, tcSplitTyConApp, GET_TYVAR, tcSplitAppTy_maybe)
import           TcSMonad hiding (tcLookupClass)
import           Type
import           TysPrim (alphaTys)
#endif

fundepPlugin :: TcPlugin
fundepPlugin = TcPlugin
  { tcPluginInit =
      (,) <$> tcPluginIO (newIORef S.empty)
          <*> polysemyStuff
  , tcPluginSolve = solveFundep
  , tcPluginStop = const $ pure ()
#if __GLASGOW_HASKELL__ >= 904
  , tcPluginRewrite = \ _ -> emptyUFM
#endif
  }


------------------------------------------------------------------------------
-- | Like 'PredType', but has an 'Ord' instance.
newtype PredType' = PredType' { getPredType :: PredType }
  deriving newtype Outputable

instance Eq PredType' where
  (==) = ((== EQ) .) . compare

instance Ord PredType' where
  compare = nonDetCmpType `on` getPredType


------------------------------------------------------------------------------
-- | Corresponds to a 'Polysemy.Internal.Union.Find' constraint. For example,
-- given @Member (State s) r@, we would get:
data FindConstraint = FindConstraint
  { fcLoc        :: CtLoc
  , fcEffectName :: Type  -- ^ @State@
  , fcEffect     :: Type  -- ^ @State s@
  , fcRow        :: Type  -- ^ @r@
  }

instance Outputable FindConstraint where
  ppr FindConstraint{..} = parens $ sep
    [ text "effect name = " <+> ppr fcEffectName
    , text "effect = " <+> ppr fcEffect
    , text "row = " <+> ppr fcRow
    ]


------------------------------------------------------------------------------
-- | Given a list of constraints, filter out the 'FindConstraint's.
getFindConstraints :: PolysemyStuff 'Things -> [Ct] -> [FindConstraint]
getFindConstraints (findClass -> cls) cts = do
  cd@CDictCan{cc_class = cls', cc_tyargs = [eff, r]} <- cts
  guard $ cls == cls'
  pure $ FindConstraint
    { fcLoc = ctLoc cd
    , fcEffectName = getEffName eff
    , fcEffect = eff
    , fcRow = r
    }


------------------------------------------------------------------------------
-- | Get evidence in scope that aren't the 'FindConstraint's.
getExtraEvidence :: PolysemyStuff 'Things -> [Ct] -> [PredType]
getExtraEvidence things cts = do
  CDictCan{cc_class = cls, cc_tyargs = as} <- cts
  guard $ cls /= findClass things
  pure $ mkAppTys (mkTyConTy $ classTyCon cls) as


------------------------------------------------------------------------------
-- | If there's a unique given @Member@ that would cause the program to
-- typecheck, use it.
findMatchingEffectIfSingular
    :: [PredType]        -- ^ Extra wanteds
    -> Set PredType'     -- ^ Extra givens
    -> FindConstraint    -- ^ Goal
    -> [FindConstraint]  -- ^ Member constraints
    -> TcM (Maybe Type)
findMatchingEffectIfSingular
    extra_wanted
    extra_given
    (FindConstraint _ eff_name wanted r)
    ts =
  let skolems = S.fromList $ foldMap (tyCoVarsOfTypeWellScoped . fcEffect) ts
      -- Which members unify with our current goal?
      results = do
        FindConstraint _ eff_name' eff' r' <- ts
        guard $ eqType eff_name eff_name'
        guard $ eqType r r'
        subst <- maybeToList $ unify (FunctionDef skolems) wanted eff'
        pure (eff', subst)
   in case results of
        [] -> pure Nothing
        -- If there is a unique member which unifies, return it.
        [(a, _)] -> pure $ Just a
        _ ->
          -- Otherwise, check if the extra wanteds give us enough information
          -- to make a unique choice.
          --
          -- For example, if we're trying to solve @Member (State a) r@, with
          -- candidates @Members (State Int, State String) r@ and can prove
          -- that @Num a@, then we can uniquely choose @State Int@.
          fmap (singleListToJust . join) $ for results $ \(eff, subst) ->
            fmap maybeToList $
              anyM (checkExtraEvidence extra_given subst) extra_wanted >>= \case
                True -> pure $ Just eff
                False -> pure Nothing

------------------------------------------------------------------------------
-- | @checkExtraEvidence givens subst c@ returns 'True' iff we can prove that
-- the constraint @c@ holds under the substitution @subst@ in the context of
-- @givens@.
checkExtraEvidence ::
  Set PredType' ->
  SUBST ->
  PredType ->
  TcM Bool
checkExtraEvidence givens subst = flip evalStateT givens . getInstance . substTy subst


------------------------------------------------------------------------------
-- | Given an effect, compute its effect name.
getEffName :: Type -> Type
getEffName t = fst $ splitAppTys t


------------------------------------------------------------------------------
-- | Generate a wanted unification for the effect described by the
-- 'FindConstraint' and the given effect.
mkWantedForce
  :: FindConstraint
  -> Type
  -> TcPluginM (Unification, Ct)
mkWantedForce fc given = do
#if __GLASGOW_HASKELL__ >= 904
  ((ev, _), _) <- unsafeTcPluginTcM
           . runTcS
           $ newWantedEq (fcLoc fc) emptyRewriterSet Nominal wanted given
#else
  (ev, _) <- unsafeTcPluginTcM
           . runTcSDeriveds
           $ newWantedEq (fcLoc fc) Nominal wanted given
#endif
  pure ( Unification (OrdType wanted) (OrdType given)
       , CNonCanonical ev
       )
  where
    wanted = fcEffect fc


------------------------------------------------------------------------------
-- | It's very important that we don't try to unify entire effects when we're
-- in interpreter mode. It's OK to unify @T x ~ T y@, but never @e ~ T y@. This
-- function takes then "given" of an interpreter, and produces a singleton
-- skolem set iff the outermost effect to be unified is a tyvar.
skolemsForInterpreter :: Type -> Set TyVar
skolemsForInterpreter ty =
  case tcSplitAppTy_maybe ty of
    Just (GET_TYVAR -> Just skolem, _) -> S.singleton skolem
    _ -> maybe mempty S.singleton $ GET_TYVAR ty


------------------------------------------------------------------------------
-- | Generate a wanted unification for the effect described by the
-- 'FindConstraint' and the given effect --- if they can be unified in this
-- context.
mkWanted
    :: FindConstraint
    -> SolveContext
    -> Type  -- ^ The given effect.
    -> TcPluginM (Maybe (Unification, Ct))
mkWanted fc solve_ctx given = do
  whenA (not (mustUnify solve_ctx) || isJust (unify solve_ctx wanted given)) $
    mkWantedForce fc given
  where
    wanted = fcEffect fc


------------------------------------------------------------------------------
-- | Determine if there is exactly one wanted find for the @r@ in question.
exactlyOneWantedForR
    :: [FindConstraint]  -- ^ Wanted finds
    -> Type              -- ^ Effect row
    -> Bool
exactlyOneWantedForR wanteds
    = fromMaybe False
    . flip M.lookup singular_r
    . OrdType
  where
    singular_r = M.fromList
               -- TODO(sandy): Nothing fails if this is just @second (const
               -- True)@. Why not? Incomplete test suite, or doing too much
               -- work?
               . fmap (second (/= 1))
               . countLength
               $ OrdType . fcRow <$> wanteds


------------------------------------------------------------------------------
-- | Returns 'True' if we can prove the given 'PredType' has a (fully
-- instantiated) instance. Uses 'StateT' to cache the results of any instances
-- it needs to prove in service of the original goal.
getInstance :: PredType -> StateT (Set PredType') TcM Bool
getInstance predty = do
  givens <- get
  case S.member (PredType' predty) givens of
    True -> pure True
    False ->
      let (con, apps) = tcSplitTyConApp predty
      in case tyConClass_maybe con of
        Nothing -> pure False
        Just cls -> do
          env <- lift tcGetInstEnvs
          let (mres, _, _) = lookupInstEnv False env cls apps
          case mres of
            ((inst, mapps) : _) -> do
              -- Get the instantiated type of the dictionary
              let df = piResultTys (idType $ is_dfun inst)
                    $ zipWith fromMaybe alphaTys mapps
              -- pull off its resulting arguments
              let (theta, _) = tcSplitPhiTy df
              allM getInstance theta >>= \case
                True -> do
                  -- Record that we can solve this instance, in case it's used
                  -- elsewhere
                  modify $ S.insert $ coerce predty
                  pure True
                False -> pure False
            _ -> pure False


solveFundep
    :: ( IORef (S.Set Unification)
       , PolysemyStuff 'Things
       )
#if __GLASGOW_HASKELL__ >= 904
    -> EvBindsVar
    -> [Ct]
    -> [Ct]
    -> TcPluginM TcPluginSolveResult
#else
    -> [Ct]
    -> [Ct]
    -> [Ct]
    -> TcPluginM TcPluginResult
#endif
solveFundep _ _ _ [] = pure $ TcPluginOk [] []
#if __GLASGOW_HASKELL__ >= 904
solveFundep (ref, stuff) _ given wanted = do
#else
solveFundep (ref, stuff) given _ wanted = do
#endif
  let wanted_finds = getFindConstraints stuff wanted
      given_finds  = getFindConstraints stuff given
      extra_wanted = getExtraEvidence stuff wanted
      extra_given = S.fromList $ coerce $ getExtraEvidence stuff given

  eqs <- forM wanted_finds $ \fc -> do
    let r  = fcRow fc
    res <- unsafeTcPluginTcM
         $ findMatchingEffectIfSingular extra_wanted extra_given fc given_finds
    case res of
      -- We found a real given, therefore we are in the context of a function
      -- with an explicit @Member e r@ constraint. We also know it can
      -- be unified (although it may generate unsatisfiable constraints).
      Just eff' -> Just <$> mkWantedForce fc eff'

      -- Otherwise, check to see if @r ~ (e ': r')@. If so, pretend we're
      -- trying to solve a given @Member e r@. But this can only happen in the
      -- context of an interpreter!
      Nothing ->
        case splitAppTys r of
          (_, [_, eff', _]) ->
            mkWanted fc
              (InterpreterUse
                (exactlyOneWantedForR wanted_finds r)
                (skolemsForInterpreter eff'))
              eff'
          _ -> pure Nothing

  -- We only want to emit a unification wanted once, otherwise a type error can
  -- force the type checker to loop forever.
  already_emitted <- tcPluginIO $ readIORef ref
  let (unifications, new_wanteds) = unzipNewWanteds already_emitted $ catMaybes eqs
  tcPluginIO $ modifyIORef ref $ S.union $ S.fromList unifications

  pure $ TcPluginOk [] new_wanteds