content: {-# LANGUAGE FlexibleContexts, MultiParamTypeClasses #-}
  {-|
  Module      : Idris.Core.Constraints
  Description : Check universe constraints.
  
  License     : BSD3
  Maintainer  : The Idris Community.
  -}
  module Idris.Core.Constraints ( ucheck ) where
  
  import Idris.Core.TT (ConstraintFC(..), Err'(..), TC(..), UConstraint(..),
                        UExp(..))
  
  import Control.Monad
  import Control.Monad.State.Strict
  import Data.List (partition)
  import qualified Data.Map.Strict as M
  import qualified Data.Set as S
  
  -- | Check that a list of universe constraints can be satisfied.
  ucheck :: S.Set ConstraintFC -> TC ()
  ucheck = void . solve 10 . S.filter (not . ignore) . dropUnused
      where
          -- TODO: remove the first ignore clause once Idris.Core.Binary:598 is dealt with
          ignore (ConstraintFC c _) | any (== Var [] (-1)) (varsIn c) = True
          ignore (ConstraintFC (ULE a b) _) = a == b
          ignore _ = False
  
  dropUnused :: S.Set ConstraintFC -> S.Set ConstraintFC
  dropUnused xs = let cs = S.toList xs
                      onlhs = countLHS M.empty cs in
                      addIfUsed S.empty onlhs cs
    where
      -- Count the number of times a variable occurs on the LHS of a constraint
      countLHS ms [] = ms
      countLHS ms (c : cs) = let lhvar = getLHS (uconstraint c)
                                 num = case M.lookup lhvar ms of
                                            Nothing -> 1
                                            Just v -> v + 1 in
                                 countLHS (M.insert lhvar num ms) cs
  
      -- Only keep a constraint if the variable on the RHS is used elsewhere
      -- on the LHS of a constraint
      addIfUsed cs' lhs [] = cs'
      addIfUsed cs' lhs (c : cs)
           = let rhvar = getRHS (uconstraint c) in
                 case M.lookup rhvar lhs of
                      Nothing -> addIfUsed cs' lhs cs
                      Just v -> addIfUsed (S.insert c cs') lhs cs
  
      getLHS (ULT x _) = x
      getLHS (ULE x _) = x
  
      getRHS (ULT _ x) = x
      getRHS (ULE _ x) = x
  
  
  data Var = Var String Int
      deriving (Eq, Ord, Show)
  
  data Domain = Domain Int Int
      deriving (Eq, Ord, Show)
  
  data SolverState =
      SolverState
          { queue       :: Queue
          , domainStore :: M.Map Var ( Domain
                                     , S.Set ConstraintFC        -- constraints that effected this variable
                                     )
          , cons_lhs    :: M.Map Var (S.Set ConstraintFC)
          , cons_rhs    :: M.Map Var (S.Set ConstraintFC)
          }
  
  data Queue = Queue [ConstraintFC] (S.Set UConstraint)
  
  
  solve :: Int -> S.Set ConstraintFC -> TC (M.Map Var Int)
  solve maxUniverseLevel ucs =
      evalStateT (propagate >> extractSolution) initSolverState
  
      where
          inpConstraints = S.toAscList ucs
  
          -- | initial solver state.
          --   the queue contains all constraints, the domain store contains the initial domains.
          initSolverState :: SolverState
          initSolverState =
              let
                  (initUnaryQueue, initQueue) = partition (\ c -> length (varsIn (uconstraint c)) == 1) inpConstraints
              in
                  SolverState
                      { queue = Queue (initUnaryQueue ++ initQueue) (S.fromList (map uconstraint (initUnaryQueue ++ initQueue)))
                      , domainStore = M.fromList
                          [ (v, (Domain 0 maxUniverseLevel, S.empty))
                          | v <- ordNub [ v
                                        | ConstraintFC c _ <- inpConstraints
                                        , v <- varsIn c
                                        ]
                          ]
                      , cons_lhs = constraintsLHS
                      , cons_rhs = constraintsRHS
                      }
  
          lhs (ULT (UVar ns x) _) = Just (Var ns x)
          lhs (ULE (UVar ns x) _) = Just (Var ns x)
          lhs _ = Nothing
  
          rhs (ULT _ (UVar ns x)) = Just (Var ns x)
          rhs (ULE _ (UVar ns x)) = Just (Var ns x)
          rhs _ = Nothing
  
          -- | a map from variables to the list of constraints the variable occurs in. (in the LHS of a constraint)
          constraintsLHS :: M.Map Var (S.Set ConstraintFC)
          constraintsLHS = M.fromListWith S.union
              [ (v, S.singleton (ConstraintFC c fc))
              | (ConstraintFC c fc) <- inpConstraints
              , let vars = varsIn c
              , length vars > 1               -- do not register unary constraints
              , v <- vars
              , lhs c == Just v
              ]
  
          -- | a map from variables to the list of constraints the variable occurs in. (in the RHS of a constraint)
          constraintsRHS :: M.Map Var (S.Set ConstraintFC)
          constraintsRHS = M.fromListWith S.union
              [ (v, S.singleton (ConstraintFC c fc))
              | (ConstraintFC c fc) <- inpConstraints
              , let vars = varsIn c
              , length vars > 1               -- do not register unary constraints
              , v <- vars
              , rhs c == Just v
              ]
  
          -- | this is where the actual work is done.
          --   dequeue the first constraint,
          --   filter domains,
          --   update domains (possibly resulting in a domain wipe out),
          --   until the queue is empty.
          propagate :: StateT SolverState TC ()
          propagate = do
              mcons <- nextConstraint
              case mcons of
                  Nothing -> return ()
                  Just (ConstraintFC cons fc) -> do
                      case cons of
                          ULE a b -> do
                              Domain lowerA upperA <- domainOf a
                              Domain lowerB upperB <- domainOf b
                              when (upperB < upperA) $ updateUpperBoundOf (ConstraintFC cons fc) a upperB
                              when (lowerA > lowerB) $ updateLowerBoundOf (ConstraintFC cons fc) b lowerA
                          ULT a b -> do
                              Domain lowerA upperA <- domainOf a
                              Domain lowerB upperB <- domainOf b
                              let upperB_pred = pred upperB
                              let lowerA_succ = succ lowerA
                              when (upperB_pred < upperA) $ updateUpperBoundOf (ConstraintFC cons fc) a upperB_pred
                              when (lowerA_succ > lowerB) $ updateLowerBoundOf (ConstraintFC cons fc) b lowerA_succ
                      propagate
  
          -- | extract a solution from the state.
          extractSolution :: (MonadState SolverState m, Functor m) => m (M.Map Var Int)
          extractSolution = M.map (extractValue . fst) <$> gets domainStore
  
          extractValue :: Domain -> Int
          extractValue (Domain x _) = x
  
          -- | dequeue the first constraint.
          nextConstraint :: MonadState SolverState m => m (Maybe ConstraintFC)
          nextConstraint = do
              Queue list set <- gets queue
              case list of
                  [] -> return Nothing
                  (q:qs) -> do
                      modify $ \ st -> st { queue = Queue qs (S.delete (uconstraint q) set) }
                      return (Just q)
  
          -- | look up the domain of a variable from the state.
          --   for convenience, this function also accepts UVal's and returns a singleton domain for them.
          domainOf :: MonadState SolverState m => UExp -> m Domain
          domainOf (UVar ns var) = gets (fst . (M.! Var ns var) . domainStore)
          domainOf (UVal val) = return (Domain val val)
  
          asPair :: Domain -> (Int, Int)
          asPair (Domain x y) = (x, y)
  
          updateUpperBoundOf :: ConstraintFC -> UExp -> Int -> StateT SolverState TC ()
          updateUpperBoundOf suspect (UVar ns var) upper = do
              doms <- gets domainStore
              let (oldDom@(Domain lower _), suspects) = doms M.! Var ns var
              let newDom = Domain lower upper
              when (wipeOut newDom) $
                lift $ Error $
                  UniverseError (ufc suspect) (UVar ns var)
                                (asPair oldDom) (asPair newDom)
                                (suspect : S.toList suspects)
              modify $ \ st -> st { domainStore = M.insert (Var ns var) (newDom, S.insert suspect suspects) doms }
              addToQueueRHS (uconstraint suspect) (Var ns var)
          updateUpperBoundOf _ UVal{} _ = return ()
  
          updateLowerBoundOf :: ConstraintFC -> UExp -> Int -> StateT SolverState TC ()
          updateLowerBoundOf suspect (UVar ns var) lower = do
              doms <- gets domainStore
              let (oldDom@(Domain _ upper), suspects) = doms M.! Var ns var
              let newDom = Domain lower upper
              when (wipeOut newDom) $
                lift $ Error $
                  UniverseError (ufc suspect) (UVar ns var)
                                (asPair oldDom) (asPair newDom)
                                (suspect : S.toList suspects)
              modify $ \ st -> st { domainStore = M.insert (Var ns var) (newDom, S.insert suspect suspects) doms }
              addToQueueLHS (uconstraint suspect) (Var ns var)
          updateLowerBoundOf _ UVal{} _ = return ()
  
          -- | add all constraints (with the given var on the lhs) to the queue
          addToQueueLHS :: MonadState SolverState m => UConstraint -> Var -> m ()
          addToQueueLHS thisCons var = do
              clhs <- gets cons_lhs
              case M.lookup var clhs of
                  Nothing -> return ()
                  Just cs -> do
                      Queue list set <- gets queue
                      let set' = S.insert thisCons set
                      let newCons = [ c | c <- S.toList cs, uconstraint c `S.notMember` set' ]
                      if null newCons
                          then return ()
                          else modify $ \ st -> st { queue = Queue (list ++ newCons)
                                                                   (S.union set (S.fromList (map uconstraint newCons))) }
  
          -- | add all constraints (with the given var on the rhs) to the queue
          addToQueueRHS :: MonadState SolverState m => UConstraint -> Var -> m ()
          addToQueueRHS thisCons var = do
              crhs <- gets cons_rhs
              case M.lookup var crhs of
                  Nothing -> return ()
                  Just cs -> do
                      Queue list set <- gets queue
                      let set' = S.insert thisCons set
                      let newCons = [ c | c <- S.toList cs, uconstraint c `S.notMember` set' ]
                      if null newCons
                          then return ()
                          else modify $ \ st -> st { queue = Queue (list ++ newCons)
                                                                   (insertAll (map uconstraint newCons) set) }
  
          insertAll [] s = s
          insertAll (x : xs) s = insertAll xs (S.insert x s)
  
          -- | check if a domain is wiped out.
          wipeOut :: Domain -> Bool
          wipeOut (Domain l u) = l > u
  
  ordNub :: Ord a => [a] -> [a]
  ordNub = S.toList . S.fromList
  
  -- | variables in a constraint
  varsIn :: UConstraint -> [Var]
  varsIn (ULT a b) = [ Var ns v | UVar ns v <- [a,b] ]
  varsIn (ULE a b) = [ Var ns v | UVar ns v <- [a,b] ]
  