module Type.Solve (solve) where

import Control.Applicative ((<$>), (<*>))
import Control.Monad
import Control.Monad.State (StateT, execStateT, liftIO)
import Control.Monad.Except (ExceptT, throwError)
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Traversable as Traversable
import qualified Data.UnionFind.IO as UF

import qualified Reporting.Annotation as A
import qualified Reporting.Error.Type as Error
import qualified Type.State as TS
import Type.Type
import Type.Unify


{-| Every variable has rank less than or equal to the maxRank of the pool.
This sorts variables into the young and old pools accordingly.
-}
generalize :: TS.Pool -> StateT TS.SolverState IO ()
generalize youngPool =
  do  youngMark <- TS.uniqueMark
      let youngRank = TS.maxRank youngPool
      let insert dict var =
            do  descriptor <- liftIO $ UF.descriptor var
                liftIO $ UF.modifyDescriptor var (\desc -> desc { mark = youngMark })
                return $ Map.insertWith (++) (rank descriptor) [var] dict

      -- Sort the youngPool variables by rank.
      rankDict <- foldM insert Map.empty (TS.inhabitants youngPool)

      -- get the ranks right for each entry.
      -- start at low ranks so that we only have to pass
      -- over the information once.
      visitedMark <- TS.uniqueMark
      forM (Map.toList rankDict) $ \(poolRank, vars) ->
          mapM (adjustRank youngMark visitedMark poolRank) vars

      -- For variables that have rank lowerer than youngRank, register them in
      -- the old pool if they are not redundant.
      let registerIfNotRedundant var =
            do  isRedundant <- liftIO $ UF.redundant var
                if isRedundant then return var else TS.register var

      let rankDict' = Map.delete youngRank rankDict
      Traversable.traverse (mapM registerIfNotRedundant) rankDict'

      -- For variables with rank youngRank
      --   If rank < youngRank: register in oldPool
      --   otherwise generalize
      let registerIfLowerRank var = do
            isRedundant <- liftIO $ UF.redundant var
            case isRedundant of
              True -> return ()
              False -> do
                desc <- liftIO $ UF.descriptor var
                case rank desc < youngRank of
                  True -> TS.register var >> return ()
                  False -> do
                    let flex' = case flex desc of { Flexible -> Rigid ; other -> other }
                    liftIO $ UF.setDescriptor var (desc { rank = noRank, flex = flex' })

      mapM_ registerIfLowerRank (Map.findWithDefault [] youngRank rankDict)


-- adjust the ranks of variables such that ranks never increase as you
-- move deeper into a variable.
adjustRank :: Int -> Int -> Int -> Variable -> StateT TS.SolverState IO Int
adjustRank youngMark visitedMark groupRank var =
  do  descriptor <- liftIO $ UF.descriptor var
      case () of
        ()  | mark descriptor == youngMark ->
                do  -- Set the variable as marked first because it may be cyclic.
                    liftIO $ UF.modifyDescriptor var $ \desc ->
                        desc { mark = visitedMark }
                    rank' <- maybe (return groupRank) adjustTerm (structure descriptor)
                    liftIO $ UF.modifyDescriptor var $ \desc ->
                        desc { rank = rank' }
                    return rank'

            | mark descriptor /= visitedMark ->
                do  let rank' = min groupRank (rank descriptor)
                    liftIO $ UF.setDescriptor var (descriptor { mark = visitedMark, rank = rank' })
                    return rank'

            | otherwise ->
                return (rank descriptor)
  where
    adjust =
        adjustRank youngMark visitedMark groupRank

    adjustTerm term =
        case term of
          App1 a b ->
              max `liftM` adjust a `ap` adjust b

          Fun1 a b ->
              max `liftM` adjust a `ap` adjust b

          Var1 x ->
              adjust x

          EmptyRecord1 ->
              return outermostRank

          Record1 fields extension ->
              do  ranks <- mapM adjust (concat (Map.elems fields))
                  rnk <- adjust extension
                  return . maximum $ rnk : ranks


-- SOLVER

solve :: TypeConstraint -> ExceptT [A.Located Error.Error] IO TS.SolverState
solve constraint =
  do  state <-
          liftIO (execStateT (actuallySolve constraint) TS.initialState)
      case TS.sError state of
        [] ->
            return state
        errors ->
            throwError errors


actuallySolve :: TypeConstraint -> StateT TS.SolverState IO ()
actuallySolve constraint =
  case constraint of
    CTrue ->
        return ()

    CSaveEnv ->
        TS.saveLocalEnv

    CEqual region term1 term2 ->
        do  t1 <- TS.flatten term1
            t2 <- TS.flatten term2
            unify region t1 t2

    CAnd cs ->
        mapM_ actuallySolve cs

    CLet [Scheme [] fqs constraint' _] CTrue ->
        do  oldEnv <- TS.getEnv
            mapM TS.introduce fqs
            actuallySolve constraint'
            TS.modifyEnv (\_ -> oldEnv)

    CLet schemes constraint' ->
        do  oldEnv <- TS.getEnv
            headers <- Map.unions `fmap` mapM solveScheme schemes
            TS.modifyEnv $ \env -> Map.union headers env
            actuallySolve constraint'
            mapM occurs $ Map.toList headers
            TS.modifyEnv (\_ -> oldEnv)

    CInstance region name term ->
        do  env <- TS.getEnv
            freshCopy <-
                case Map.lookup name env of
                  Just tipe ->
                      TS.makeInstance tipe
                  Nothing ->
                      if List.isPrefixOf "Native." name
                        then liftIO (variable Flexible)
                        else error ("Could not find '" ++ name ++ "' when solving type constraints.")

            t <- TS.flatten term
            unify region freshCopy t


solveScheme :: TypeScheme -> StateT TS.SolverState IO (Map.Map String Variable)
solveScheme scheme =
  case scheme of
    Scheme [] [] constraint header ->
        do  actuallySolve constraint
            Traversable.traverse TS.flatten header

    Scheme rigidQuantifiers flexibleQuantifiers constraint header ->
        do  let quantifiers = rigidQuantifiers ++ flexibleQuantifiers
            oldPool <- TS.getPool

            -- fill in a new pool when working on this scheme's constraints
            freshPool <- TS.nextRankPool
            TS.switchToPool freshPool
            mapM TS.introduce quantifiers
            header' <- Traversable.traverse TS.flatten header
            actuallySolve constraint

            allDistinct rigidQuantifiers
            youngPool <- TS.getPool
            TS.switchToPool oldPool
            generalize youngPool
            mapM isGeneric rigidQuantifiers
            return header'


-- ADDITIONAL CHECKS

-- Checks that all of the given variables belong to distinct equivalence classes.
-- Also checks that their structure is Nothing, so they represent a variable, not
-- a more complex term.
allDistinct :: [Variable] -> StateT TS.SolverState IO ()
allDistinct vars =
  do  seen <- TS.uniqueMark
      forM_ vars $ \var ->
        do  desc <- liftIO $ UF.descriptor var
            case structure desc of
              Just _ ->
                TS.addError (error "Cannot generalize something that is not a type variable.") undefined

              Nothing ->
                do  when (mark desc == seen) $
                      TS.addError (error "Duplicate variable during generalization.") undefined
                    liftIO $ UF.setDescriptor var (desc { mark = seen })


-- Check that a variable has rank == noRank, meaning that it can be generalized.
isGeneric :: Variable -> StateT TS.SolverState IO ()
isGeneric var =
  do  desc <- liftIO $ UF.descriptor var
      if rank desc == noRank
        then return ()
        else
          TS.addError (error "Unable to generalize a type variable. It is not unranked.") undefined


occurs :: (String, Variable) -> StateT TS.SolverState IO ()
occurs (_name, variable) =
  do  vars <- liftIO $ infiniteVars [] variable
      case vars of
        [] ->
          return ()

        var : _ ->
          do  desc <- liftIO $ UF.descriptor var
              case structure desc of
                Nothing ->
                  error "occurs check triggered"

                Just _ ->
                  error "occurs check triggered"
  where
    infiniteVars :: [Variable] -> Variable -> IO [Variable]
    infiniteVars seen var =
      let go = infiniteVars (var:seen)
      in
      if var `elem` seen
        then return [var]
        else do
          desc <- UF.descriptor var
          case structure desc of
            Nothing -> return []
            Just struct ->
                case struct of
                  App1 a b ->
                      (++) <$> go a <*> go b

                  Fun1 a b ->
                      (++) <$> go a <*> go b

                  Var1 a ->
                      go a

                  EmptyRecord1 ->
                      return []

                  Record1 fields ext ->
                      concat <$> mapM go (ext : concat (Map.elems fields))