<commit_before>module Type.Solve (solve) where

import Control.Monad
import Control.Monad.State
import qualified Data.UnionFind.IO as UF
import qualified Data.Map as Map
import qualified Data.Traversable as Traversable
import qualified Data.Maybe as Maybe
import qualified Data.List as List
import Type.Type
import Type.Unify
import qualified Type.Environment as Env
import qualified Type.State as TS
import qualified Text.PrettyPrint as P
import SourceSyntax.Location (Located(L), SrcSpan)


-- | 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
      insert dict var = do
        desc <- liftIO $ UF.descriptor var
        liftIO $ UF.modifyDescriptor var (\desc -> desc { mark = youngMark })
        return $ Map.insertWith (++) (rank desc) [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
  mapM (\(poolRank, vars) -> mapM (adjustRank youngMark visitedMark poolRank) vars) (Map.toList rankDict)

  -- 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
        if isRedundant then return () else do
            desc <- liftIO $ UF.descriptor var
            if rank desc < youngRank
              then TS.register var >> return ()
              else let flex' = if flex desc == Flexible then Rigid else flex desc
                   in  liftIO $ UF.setDescriptor var (desc { rank = noRank, flex = flex' })

  mapM registerIfLowerRank (Map.findWithDefault [] youngRank rankDict)

  return ()


-- 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 variable =
    let adjust = adjustRank youngMark visitedMark groupRank in
    do desc <- liftIO $ UF.descriptor variable
       case () of
         () | mark desc == youngMark ->
                do rank' <- case structure desc of
                              Nothing -> return groupRank
                              Just 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
                   liftIO $ UF.setDescriptor variable (desc { mark = visitedMark, rank = rank' })
                   return rank'

            | mark desc /= visitedMark ->
                do let rank' = min groupRank (rank desc)<commit_msg>style tweaks<commit_after>module Type.Solve (solve) where

import Control.Monad
import Control.Monad.State
import qualified Data.UnionFind.IO as UF
import qualified Data.Map as Map
import qualified Data.Traversable as Traversable
import qualified Data.Maybe as Maybe
import qualified Data.List as List
import Type.Type
import Type.Unify
import qualified Type.Environment as Env
import qualified Type.State as TS
import qualified Text.PrettyPrint as P
import SourceSyntax.Location (Located(L), SrcSpan)


-- | 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
      insert dict var = do
        desc <- liftIO $ UF.descriptor var
        liftIO $ UF.modifyDescriptor var (\desc -> desc { mark = youngMark })
        return $ Map.insertWith (++) (rank desc) [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
  mapM (\(poolRank, vars) -> mapM (adjustRank youngMark visitedMark poolRank) vars) (Map.toList rankDict)

  -- 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 variable =
    let adjust = adjustRank youngMark visitedMark groupRank in
    do desc <- liftIO $ UF.descriptor variable
       case () of
         () | mark desc == youngMark ->
                do rank' <- case structure desc of
                              Nothing -> return groupRank
                              Just 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
                   liftIO $ UF.setDescriptor variable (desc { mark = visitedMark, rank = rank' })
                   return rank'

            | mark desc /= visitedMark ->
                do let rank' = min groupRank (rank desc)
{'fraction_non_alphanumeric': 0.06831734508685507, 'fraction_numerical': 0.0012963443090484833, 'mean_word_length': 1.8971085242208037, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4515960', 'n_tokens_mistral': 1993, 'n_tokens_neox': 1883, 'n_words': 796}