<commit_before>{-# LANGUAGE BangPatterns #-}
module Distribution.Solver.Modular.Explore
    ( backjump
    , backjumpAndExplore
    ) where

import Control.Monad.State.Lazy
import Data.Foldable as F
import Data.List as L (foldl')
import Data.Map as M

import Distribution.Solver.Modular.Assignment
import Distribution.Solver.Modular.Dependency
import Distribution.Solver.Modular.Log
import Distribution.Solver.Modular.Message
import qualified Distribution.Solver.Modular.PSQ as P
import qualified Distribution.Solver.Modular.ConflictSet as CS
import Distribution.Solver.Modular.Tree
import Distribution.Solver.Types.PackagePath
import Distribution.Solver.Types.Settings (EnableBackjumping(..), CountConflicts(..))
import qualified Distribution.Solver.Types.Progress as P

type Explore = State ConflictMap

-- | This function takes the variable we're currently considering, an
-- initial conflict set and a
-- list of children's logs. Each log yields either a solution or a
-- conflict set. The result is a combined log for the parent node that
-- has explored a prefix of the children.
--
-- We can stop traversing the children's logs if we find an individual
-- conflict set that does not contain the current variable. In this
-- case, we can just lift the conflict set to the current level,
-- because the current level cannot possibly have contributed to this
-- conflict, so no other choice at the current level would avoid the
-- conflict.
--
-- If any of the children might contain a successful solution, we can
-- return it immediately. If all children contain conflict sets, we can
-- take the union as the combined conflict set.
--
-- The initial conflict set corresponds to the justification that we
-- have to choose this goal at all. There is a reason why we have
-- introduced the goal in the first place, and this reason is in conflict
-- with the (virtual) option not to choose anything for the current
-- variable. See also the comments for 'avoidSet'.
--
backjump :: EnableBackjumping -> Var QPN
         -> ConflictSet QPN -> P.PSQ k (Explore (ConflictSetLog a))
         -> Explore (ConflictSetLog a)
backjump (EnableBackjumping enableBj) var initial xs =
    makeStrict $ F.foldr combine logBackjump xs initial
  where
    combine ::  Explore (ConflictSetLog a)
            -> (ConflictSet QPN -> Explore (ConflictSetLog a))
            ->  ConflictSet QPN -> Explore (ConflictSetLog a)
    combine x f csAcc = do
      l <- x
      case l of
        P.Done d  -> return (P.Done d)
        P.Fail cs
          | enableBj && not (var `CS.member` cs) -> logBackjump cs
          | otherwise                            -> f (csAcc `CS.union` cs)
        P.Step m ms -> do
          l' <- combine (return ms) f csAcc
          return (P.Step m l')

    logBackjump :: ConflictSet QPN -> Explore (ConflictSetLog a)
    logBackjump cs = return (failWith (Failure cs Backjump) cs)

type ConflictSetLog = P.Progress Message (ConflictSet QPN)

type ConflictMap = Map (Var QPN) Int

getBestGoal :: ConflictMap -> P.PSQ (Goal QPN) a -> (Goal QPN, a)
getBestGoal cm =
  P.maximumBy
    ( flip (M.findWithDefault 0) cm
    . (\ (Goal v _) -> v)
    )

getFirstGoal :: P.PSQ (Goal QPN) a -> (Goal QPN, a)
getFirstGoal ts =
  P.casePSQ ts
    (error "getFirstGoal: empty goal choice") -- empty goal choice is an internal error
    (\ k v _xs -> (k, v))  -- commit to the first goal choice

updateCM :: ConflictSet QPN -> ConflictMap -> ConflictMap
updateCM cs cm =
  L.foldl' (\ cmc k -> M.alter inc k cmc) cm (CS.toList cs)
  where
    inc Nothing  = Just 1
    inc (Just n) = Just $! n + 1

makeStrict :: Explore a -> Explore a
makeStrict e = state (\ !cm -> runState e cm)

-- | A tree traversal that simultaneously propagates conflict sets up
-- the tree from the leaves and creates a log.
exploreLog :: EnableBackjumping -> CountConflicts -> Tree QGoalReason
           -> (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
exploreLog enableBj (CountConflicts countConflicts) = cata (\ x y -> makeStrict (go x y))
  where
    updateCM' :: ConflictSet QPN -> Explore ()
    updateCM'
      | countConflicts = modify' . updateCM
      | otherwise      = const (modify' id)

    getBestGoal' :: P.PSQ (Goal QPN) a -> Explore (Goal QPN, a)
    getBestGoal'
      | countConflicts = \ ts -> get >>= \ cm -> return (getBestGoal cm ts)
      | otherwise      = return . getFirstGoal

    go :: TreeF QGoalReason (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
                         -> (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
    go (FailF c fr)             _            = updateCM' c >> return (failWith (Failure c fr) c)
    go (DoneF rdm)              a            = return (succeedWith Success (a, rdm))
    go (PChoiceF qpn gr     ts) (A pa fa sa) =
      backjump enableBj (P qpn) (avoidSet (P qpn) gr) $ -- try children in order,
        P.mapWithKey                                -- when descending ...
          (\ i@(POption k _) r -> do
            l <- r (A (M.insert qpn k pa) fa sa)
            return (tryWith (TryP qpn i) l)
          )
        ts
    go (FChoiceF qfn gr _ _ ts) (A pa fa sa) =
      backjump enableBj (F qfn) (avoidSet (F qfn) gr) $ -- try children in order,
        P.mapWithKey                                -- when descending ...
          (\ k r -> do
            l <- r (A pa (M.insert qfn k fa) sa)
            return (tryWith (TryF qfn k) l)
          )
        ts
    go (SChoiceF qsn gr _   ts) (A pa fa sa) =<commit_msg>Better way to prevent space leak for --no-count-conflicts.<commit_after>{-# LANGUAGE BangPatterns #-}
module Distribution.Solver.Modular.Explore
    ( backjump
    , backjumpAndExplore
    ) where

import Control.Monad.State.Lazy
import Data.Foldable as F
import Data.List as L (foldl')
import Data.Map as M

import Distribution.Solver.Modular.Assignment
import Distribution.Solver.Modular.Dependency
import Distribution.Solver.Modular.Log
import Distribution.Solver.Modular.Message
import qualified Distribution.Solver.Modular.PSQ as P
import qualified Distribution.Solver.Modular.ConflictSet as CS
import Distribution.Solver.Modular.Tree
import Distribution.Solver.Types.PackagePath
import Distribution.Solver.Types.Settings (EnableBackjumping(..), CountConflicts(..))
import qualified Distribution.Solver.Types.Progress as P

type Explore = State ConflictMap

-- | This function takes the variable we're currently considering, an
-- initial conflict set and a
-- list of children's logs. Each log yields either a solution or a
-- conflict set. The result is a combined log for the parent node that
-- has explored a prefix of the children.
--
-- We can stop traversing the children's logs if we find an individual
-- conflict set that does not contain the current variable. In this
-- case, we can just lift the conflict set to the current level,
-- because the current level cannot possibly have contributed to this
-- conflict, so no other choice at the current level would avoid the
-- conflict.
--
-- If any of the children might contain a successful solution, we can
-- return it immediately. If all children contain conflict sets, we can
-- take the union as the combined conflict set.
--
-- The initial conflict set corresponds to the justification that we
-- have to choose this goal at all. There is a reason why we have
-- introduced the goal in the first place, and this reason is in conflict
-- with the (virtual) option not to choose anything for the current
-- variable. See also the comments for 'avoidSet'.
--
backjump :: EnableBackjumping -> Var QPN
         -> ConflictSet QPN -> P.PSQ k (Explore (ConflictSetLog a))
         -> Explore (ConflictSetLog a)
backjump (EnableBackjumping enableBj) var initial xs =
    F.foldr combine logBackjump xs initial
  where
    combine ::  Explore (ConflictSetLog a)
            -> (ConflictSet QPN -> Explore (ConflictSetLog a))
            ->  ConflictSet QPN -> Explore (ConflictSetLog a)
    combine x f csAcc = do
      l <- x
      case l of
        P.Done d  -> return (P.Done d)
        P.Fail cs
          | enableBj && not (var `CS.member` cs) -> logBackjump cs
          | otherwise                            -> f (csAcc `CS.union` cs)
        P.Step m ms -> do
          l' <- combine (return ms) f csAcc
          return (P.Step m l')

    logBackjump :: ConflictSet QPN -> Explore (ConflictSetLog a)
    logBackjump cs = return (failWith (Failure cs Backjump) cs)

type ConflictSetLog = P.Progress Message (ConflictSet QPN)

type ConflictMap = Map (Var QPN) Int

getBestGoal :: ConflictMap -> P.PSQ (Goal QPN) a -> (Goal QPN, a)
getBestGoal cm =
  P.maximumBy
    ( flip (M.findWithDefault 0) cm
    . (\ (Goal v _) -> v)
    )

getFirstGoal :: P.PSQ (Goal QPN) a -> (Goal QPN, a)
getFirstGoal ts =
  P.casePSQ ts
    (error "getFirstGoal: empty goal choice") -- empty goal choice is an internal error
    (\ k v _xs -> (k, v))  -- commit to the first goal choice

updateCM :: ConflictSet QPN -> ConflictMap -> ConflictMap
updateCM cs cm =
  L.foldl' (\ cmc k -> M.alter inc k cmc) cm (CS.toList cs)
  where
    inc Nothing  = Just 1
    inc (Just n) = Just $! n + 1

-- | A tree traversal that simultaneously propagates conflict sets up
-- the tree from the leaves and creates a log.
exploreLog :: EnableBackjumping -> CountConflicts -> Tree QGoalReason
           -> (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
exploreLog enableBj (CountConflicts countConflicts) = cata go
  where
    updateCM' :: ConflictSet QPN -> Explore a -> Explore a
    updateCM'
      | countConflicts = \ c k -> modify' (updateCM c) >> k
      | otherwise      = \ _ k -> k

    getBestGoal' :: P.PSQ (Goal QPN) a -> Explore (Goal QPN, a)
    getBestGoal'
      | countConflicts = \ ts -> get >>= \ cm -> return (getBestGoal cm ts)
      | otherwise      = return . getFirstGoal

    go :: TreeF QGoalReason (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
                         -> (Assignment -> Explore (ConflictSetLog (Assignment, RevDepMap)))
    go (FailF c fr)             _            = updateCM' c (return (failWith (Failure c fr) c))
    go (DoneF rdm)              a            = return (succeedWith Success (a, rdm))
    go (PChoiceF qpn gr     ts) (A pa fa sa) =
      backjump enableBj (P qpn) (avoidSet (P qpn) gr) $ -- try children in order,
        P.mapWithKey                                -- when descending ...
          (\ i@(POption k _) r -> do
            l <- r (A (M.insert qpn k pa) fa sa)
            return (tryWith (TryP qpn i) l)
          )
        ts
    go (FChoiceF qfn gr _ _ ts) (A pa fa sa) =
      backjump enableBj (F qfn) (avoidSet (F qfn) gr) $ -- try children in order,
        P.mapWithKey                                -- when descending ...
          (\ k r -> do
            l <- r (A pa (M.insert qfn k fa) sa)
            return (tryWith (TryF qfn k) l)
          )
        ts
    go (SChoiceF qsn gr _   ts) (A pa fa sa) =
{'fraction_non_alphanumeric': 0.09515146002001273, 'fraction_numerical': 0.0005458018739197671, 'mean_word_length': 2.4453149482920713, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 85, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7054129', 'n_tokens_mistral': 3374, 'n_tokens_neox': 3249, 'n_words': 1309}