{-# LANGUAGE FlexibleInstances #-}

module Intensional.InferM
  ( InferM,
    Context,
    InferEnv (..),
    Stats (..),
    runInferM,
    Intensional.InferM.saturate,
    branchAny,
    emitDD,
    emitDK,
    emitKD,
    fresh,
    putVar,
    putVars,
    setLoc,
    getExternalName,
    isTrivial,
    isIneligible,
    noteD,
    noteK,
    incrN,
    noteErrs,
    Intensional.InferM.cexs
  )
where

import Intensional.Constraints as Constraints
import Intensional.Constructors
import Control.Monad.RWS.Strict hiding (guard)
import qualified Data.IntSet as IntSet
import qualified Data.Map as M
import GhcPlugins hiding ((<>), singleton)
import Intensional.Scheme
import Intensional.Types
import Intensional.Ubiq
import Intensional.Guard

type InferM = RWS InferEnv ConstraintSet InferState

type Context = M.Map Name Scheme

data InferEnv
  = InferEnv
      { modName :: Module, -- The current module
        varEnv :: Context,
        inferLoc :: SrcSpan -- The current location in the source text
      }

data InferState = 
        InferState {
          maxK :: Int,
          maxD :: Int,
          maxI :: Int,
          cntN :: Int,
          rVar :: Int,
          errs :: ConstraintSet
        }

initState :: InferState
initState = InferState 0 0 0 0 0 mempty

noteK :: Int -> InferM ()
noteK x = modify (\s -> s { maxK = max x (maxK s) })

noteD :: Int -> InferM ()
noteD x = modify (\s -> s { maxD = max x (maxD s) })

noteI :: Int -> InferM ()
noteI x = modify (\s -> s { maxI = max x (maxI s) })

incrV :: InferM ()
incrV = modify (\s -> s { rVar = rVar s + 1 })

incrN :: InferM ()
incrN = modify (\s -> s { cntN = cntN s + 1 })

{-|
  Given a set of trivially unsatisfiable constraints @es@,
  @noteErrs es@ is the action that records them
  in the accumulating set in the inference state.
-}
noteErrs :: ConstraintSet -> InferM ()
noteErrs es = modify (\s -> s { errs = es <> errs s })

data Stats = 
        Stats {
          getK :: Int,
          getD :: Int,
          getV :: Int,
          getI :: Int,
          getN :: Int
        }

runInferM ::
  InferM a ->
  Module ->
  Context ->
  (a, [Atomic], Stats)
runInferM run mod_name init_env =
  let (a, s, _) = runRWS run (InferEnv mod_name init_env (UnhelpfulSpan (mkFastString "Nowhere"))) initState
  in (a, Constraints.toList (errs s), Stats (maxK s) (maxD s) (rVar s) (maxI s) (cntN s))

-- Transitively remove local constraints
saturate :: Refined a => InferM a -> InferM a
saturate ma = pass $
  do
    a <- ma
    env <- asks varEnv
    m <- asks modName
    src <- asks inferLoc
    let interface = domain a <> domain env
    noteI (IntSet.size interface)
    let fn cs =
          let ds = Constraints.saturate (CInfo m src) interface cs
           in if debugging then debugBracket a env src cs ds else ds
    return (a, fn)
  where
    debugBracket a env src cs ds =
      let asz = "type: " ++ show (IntSet.size $ domain a)
          esz = "env: " ++ show (IntSet.size $ domain env)
          csz = show (size cs)
          spn = traceSpan src
          tmsg = "#interface = (" ++ asz ++ " + " ++ esz ++ "), #constraints = " ++ csz
          ds' = trace ("[TRACE] BEGIN saturate at " ++ spn ++ ": " ++ tmsg) ds
       in ds' `seq` trace ("[TRACE] END saturate at " ++ spn ++ " saturated size: " ++ (show $ size ds)) ds

{-|
    Given a constraint set @cs@, @cexs cs@ is the inference action that
    attempts to build a model of @cs@ and returns the set of counterexamples.
-}
cexs :: ConstraintSet -> InferM ConstraintSet
cexs cs = 
  do  m <- asks modName 
      src <- asks inferLoc 
      return $ Constraints.cexs (CInfo m src) cs

-- Check if a core datatype is ineligible for refinement
isIneligible :: TyCon -> InferM Bool
isIneligible tc =  
  do  m <- asks modName
      return (not (homeOrBase m (getName tc)) || null (tyConDataCons tc))
  where
    homeOrBase m n =
      nameIsHomePackage m n 

        -- Previously we tried to include as much of base as possible by asking for specific modules
        -- but this is a little too coarse grain (e.g. GHC.Types will include Bool, but also Int):
        --
        -- vc|| ( not (nameIsFromExternalPackage baseUnitId n && nameIsFromExternalPackage primUnitId n)
        --        && case nameModule_maybe n of
        --          Nothing -> False
        --          Just (Module _ m) ->
        --            List.isPrefixOf "Prelude" (moduleNameString m)
        --              || List.isPrefixOf "Data" (moduleNameString m)
        --              || List.isPrefixOf "GHC.Base" (moduleNameString m)
        --              || List.isPrefixOf "GHC.Types" (moduleNameString m)
        --    )
        --
        -- a better approach would be, simply:
        --  strName = occNameString (getOccName tc)
        --  ...
        --  || strName == "Bool"
        --  || strName == "Maybe"

isTrivial :: TyCon -> Bool
isTrivial tc = (== 1) (length (tyConDataCons tc))



{- 
  Given a list of data constructors @ks@, a datatype @Inj x d@ and an 
  inference action @m@, @branchAny ks (Inj x d) m@ is the inference action
  that consists of doing @m@ then guarding all emitted constraints
  by the requirement that @ks in x(d)@.
-}
branchAny :: [DataCon] -> DataType TyCon -> InferM a -> InferM a
branchAny _ (Base _)  m = m
branchAny ks (Inj x d) m = 
    if (isTrivial d) then m else censor guardWithAll m
  where
    dn = getName d
    guardWithAll cs =
      foldMap (\k -> Constraints.guardWith (singleton [getName k] x dn) cs) ks

mkConFromCtx :: ConL -> ConR -> InferM Atomic
mkConFromCtx l r =
  do  m <- asks modName
      s <- asks inferLoc
      return (Constraint l r mempty (CInfo m s))

emitDD :: DataType TyCon -> DataType TyCon -> InferM ()
emitDD (Inj x d) (Inj y _) = 
  unless (isTrivial d) $ 
    do  a <- mkConFromCtx (Dom (Inj x dn)) (Dom (Inj y dn))
        tell (Constraints.fromList [a])
  where
    dn = getName d
emitDD _ _ = return ()

emitKD :: DataCon -> SrcSpan -> DataType TyCon -> InferM ()
emitKD k s (Inj x d) =
  unless (isTrivial d) $ 
    do  a <- mkConFromCtx (Con kn s) (Dom (Inj x dn))
        tell (Constraints.fromList [a])
  where
    dn  = getName d
    kn = getName k
emitKD _ _ _ = return ()

emitDK :: DataType TyCon -> [DataCon] -> SrcSpan -> InferM ()
emitDK (Inj x d) ks s =
  unless (isTrivial d || length (tyConDataCons d) == length ks) $ 
    do  a <- mkConFromCtx (Dom (Inj x dn)) (Set ksn s)
        tell (Constraints.fromList [a])
  where
    dn  = getName d
    ksn = mkUniqSet (map getName ks)
emitDK _ _ _ = return ()


-- A fresh refinement variable
fresh :: InferM RVar
fresh = do
  i <- gets rVar
  incrV
  return i

-- Insert variables into environment
putVar :: Name -> Scheme -> InferM a -> InferM a
putVar n s = local (\env -> env {varEnv = M.insert n s (varEnv env)})

putVars :: Context -> InferM a -> InferM a
putVars ctx = local (\env -> env {varEnv = M.union ctx (varEnv env)})

-- Add source text location tick
setLoc :: SrcSpan -> InferM a -> InferM a
setLoc l = local (\env -> env {inferLoc = l})

-- Prepare name for interface
-- Should be used before all type variables
getExternalName :: NamedThing a => a -> InferM Name
getExternalName a = do
  let n = getName a
  mn <- asks modName
  return $ mkExternalName (nameUnique n) mn (nameOccName n) (nameSrcSpan n)