advent16/MainSPar.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
   2 -- Follow up at https://work.njae.me.uk/2023/07/21/optimising-haskell-example-3/
   3
   4 import Debug.Trace
   5
   6 import AoC
   7 import Data.Text (Text)
   8 import qualified Data.Text.IO as TIO
   9 import Data.Attoparsec.Text hiding (take, D)
  10 import Control.Applicative
  11 import qualified Data.PQueue.Prio.Max as P
  12 import qualified Data.Set as S
  13 import qualified Data.Sequence as Q
  14 import qualified Data.Map.Strict as M
  15 import Data.Map.Strict ((!))
  16 -- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>)))
  17 import Data.Sequence ( (|>), Seq((:|>)) )
  18 import Data.List
  19 import Data.List.Split (chunksOf)
  20 import Data.Ord
  21 import Control.Monad.Reader
  22 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  23 import Control.Parallel.Strategies
  24
  25
  26 type RoomID = String
  27
  28 data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
  29     deriving (Eq, Show, Ord)
  30 makeLenses ''Tunnel
  31
  32 data Room = Room
  33     { _flowRate :: Int
  34     , _tunnels :: S.Set Tunnel
  35     } deriving (Eq, Show, Ord)
  36 makeLenses ''Room
  37
  38 type Cave = M.Map RoomID Room
  39 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
  40
  41 type CaveContext = Reader TimedCave
  42
  43 data SearchState = SearchState
  44     { _currentRoom :: RoomID
  45     , _currentTime :: Int
  46     , _openValves :: S.Set RoomID
  47     } deriving (Eq, Show, Ord)
  48 makeLenses ''SearchState
  49
  50 data Agendum =
  51     Agendum { _current :: SearchState
  52             , _trail :: Q.Seq SearchState
  53             , _trailBenefit :: Int
  54             , _benefit :: Int
  55             } deriving (Show, Eq, Ord)
  56 makeLenses ''Agendum
  57
  58 type Agenda = P.MaxPQueue Int Agendum
  59
  60 -- state, total flowed so far
  61 type ExploredStates = S.Set (SearchState, Int)
  62
  63 type PartSolutions = M.Map (S.Set RoomID) Int
  64
  65
  66 main :: IO ()
  67 main =
  68   do  dataFileName <- getDataFileName
  69       text <- TIO.readFile dataFileName
  70       let expandedCave = successfulParse text
  71       -- print cave
  72       -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
  73       -- print $ compress cave
  74       -- putStrLn $ dotify expandedCave
  75       let cave = compress expandedCave
  76       print $ part1 cave
  77       print $ part2 cave
  78
  79 -- dotify cave = "graph G {\n" ++ (unlines $ concat $ M.elems $ M.mapWithKey showCRoom cave) ++ "\n}\n"
  80 --   where showCRoom roomID room = filter (not . null) ((showCRoomShape roomID room) : (showCRoomLinks roomID room))
  81
  82 -- showCRoomShape roomID room
  83 --   | room ^. flowRate > 0 = roomID ++ " [fillcolor=grey label=\"" ++ roomID ++ ": " ++ (show $ room ^. flowRate) ++ "\"];"
  84 --   | otherwise = ""
  85
  86 -- showCRoomLinks roomID room = [roomID ++ " -- " ++ (t ^. tunnelTo) ++ ";" | t <- S.toList $ room ^. tunnels, (t ^. tunnelTo) > roomID ]
  87
  88 part1, part2 :: Cave -> Int
  89 -- part1 :: Cave -> Int
  90 part1 cave = runSearch 30 cave
  91 part2 cave = maximum (fmap maximum chunkSolns `using` parList rdeepseq)
  92   where rawSolutions = runSearchAll 26 cave
  93         solutionList = M.toList rawSolutions
  94         combinations = [ fp + fe
  95                        | (p, fp) <- solutionList
  96                        , (e, fe) <- solutionList
  97                        , p < e
  98                        , S.disjoint p e
  99                        ]
 100         chunkSolns = chunksOf 10000 combinations
 101
 102 includeAgendum :: PartSolutions -> Agendum -> CaveContext PartSolutions
 103 includeAgendum results agendum =
 104   do cf <- currentFlow (agendum ^. current)
 105      timeLimit <- asks getTimeLimit
 106      let timeLeft = timeLimit - timeOf (agendum ^. current)
 107      let remainingFlow = cf * timeLeft
 108      let totalFlow = remainingFlow + agendum ^. trailBenefit
 109      let visitedSet = agendum ^. current . openValves
 110      let currentBest = M.findWithDefault 0 visitedSet results
 111      if totalFlow > currentBest
 112      then return (M.insert visitedSet totalFlow results)
 113      else return results
 114
 115 runSearch :: Int -> Cave -> Int
 116 runSearch timeLimit cave = maybe 0 _benefit result
 117     where result = runReader (searchCave "AA") (TimedCave cave timeLimit sortedRooms)
 118           sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
 119
 120 runSearchAll :: Int -> Cave -> PartSolutions
 121 runSearchAll timeLimit cave = result
 122     where result = runReader (searchCaveAll "AA") (TimedCave cave timeLimit sortedRooms)
 123           sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
 124
 125
 126 searchCave :: String -> CaveContext (Maybe Agendum)
 127 searchCave startRoom =
 128     do agenda <- initAgenda startRoom
 129        aStar agenda S.empty
 130
 131 searchCaveAll :: String -> CaveContext PartSolutions
 132 searchCaveAll startRoom =
 133     do agenda <- initAgenda startRoom
 134        allSolutions agenda S.empty M.empty
 135
 136 initAgenda :: String -> CaveContext Agenda
 137 initAgenda startID =
 138     do let startState = emptySearchState startID
 139        b <- estimateBenefit startState
 140        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 141
 142 aStar ::  Agenda -> ExploredStates -> CaveContext (Maybe Agendum)
 143 aStar agenda closed
 144     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 145     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
 146     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
 147     -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
 148     | P.null agenda = return Nothing
 149     | otherwise =
 150         do  let (_, currentAgendum) = P.findMax agenda
 151             let reached = currentAgendum ^. current
 152             nexts <- candidates currentAgendum closed
 153             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 154             reachedGoal <- isGoal currentAgendum
 155             let cl = (reached, currentAgendum ^. trailBenefit)
 156             if reachedGoal
 157             then return (Just currentAgendum)
 158             else if (cl `S.member` closed)
 159                  then aStar (P.deleteMax agenda) closed
 160                  else aStar newAgenda (S.insert cl closed)
 161
 162 allSolutions ::  Agenda -> ExploredStates -> PartSolutions -> CaveContext PartSolutions
 163 allSolutions agenda closed foundSolutions
 164     | P.null agenda = return foundSolutions
 165     | otherwise =
 166         do  let (_, currentAgendum) = P.findMax agenda
 167             let reached = currentAgendum ^. current
 168             nexts <- candidates currentAgendum closed
 169             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 170             reachedGoal <- isGoal currentAgendum
 171             let cl = (reached, currentAgendum ^. trailBenefit)
 172             newFoundSolutions <- includeAgendum foundSolutions currentAgendum
 173             if reachedGoal
 174             then allSolutions (P.deleteMax agenda) closed newFoundSolutions
 175             else if (cl `S.member` closed)
 176                  then allSolutions (P.deleteMax agenda) closed foundSolutions
 177                  else allSolutions newAgenda (S.insert cl closed) newFoundSolutions
 178
 179
 180 candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
 181 candidates agendum closed =
 182     do  let candidate = agendum ^. current
 183         let previous = agendum ^. trail
 184         let prevBenefit = agendum ^. trailBenefit
 185         succs <- successors candidate
 186         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 187         let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
 188         return nonloops
 189
 190 emptySearchState :: RoomID -> SearchState
 191 emptySearchState startID = SearchState
 192     { _currentRoom = startID
 193     , _currentTime = 0
 194     , _openValves = S.empty
 195     }
 196
 197 currentFlow :: SearchState -> CaveContext Int
 198 currentFlow state =
 199     do cave <- asks getCave
 200        let valves = state ^. openValves
 201        let presentRooms = cave `M.restrictKeys` valves
 202        return $ sumOf (folded . flowRate) presentRooms
 203
 204 timeOf :: SearchState -> Int
 205 timeOf state = state ^. currentTime
 206
 207 successors :: SearchState -> CaveContext (Q.Seq SearchState)
 208 successors state =
 209     do isFF <- isFullFlow state
 210        cave <- asks getCave
 211        timeLimit <- asks getTimeLimit
 212        let here = state ^. currentRoom
 213        let opened = state ^. openValves
 214        let now = state ^. currentTime
 215        let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
 216        let moves = [ SearchState
 217                          { _currentRoom = (t ^. tunnelTo)
 218                          , _currentTime = now + (t ^. tunnelLength)
 219                          , _openValves = opened
 220                          }
 221                    | t <- remaining
 222                    , now + (t ^. tunnelLength) <= timeLimit
 223                    ]
 224        let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 && now < timeLimit
 225                    then [SearchState { _currentRoom = here, _currentTime = now + 1, _openValves = S.insert here opened }]
 226                    else []
 227        let nexts = if null opens then moves else opens
 228        let nexts' =  if null nexts
 229                      then [ SearchState
 230                              { _currentRoom = here
 231                              , _currentTime = timeLimit
 232                              , _openValves = opened
 233                              } ]
 234                      else nexts
 235        let succs = Q.fromList nexts'
 236        if isFF || (Q.null succs)
 237           then return $ Q.singleton (state & currentTime .~ timeLimit)
 238           else return succs
 239
 240
 241 estimateBenefit :: SearchState -> CaveContext Int
 242 estimateBenefit here =
 243     do cave <- asks getCave
 244        timeLimit <- asks getTimeLimit
 245        let timeRemaining = timeLimit - (timeOf here)
 246        cf <- currentFlow here
 247        sortedValves <- asks getSortedRooms
 248        let opened = here ^. openValves
 249        let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
 250        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 251        return $ (cf * timeRemaining) + otherValveFlows
 252
 253 makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> CaveContext Agendum
 254 makeAgendum previous prevBenefit newState =
 255     do predicted <- estimateBenefit newState --  (Q.length previous)
 256        -- cf <- currentFlow newState
 257        oldFlow <- lastFlow previous (timeOf newState)
 258        let newTrail = previous |> newState
 259        let incurred = prevBenefit + oldFlow
 260        return Agendum { _current = newState
 261                       , _trail = newTrail
 262                       , _trailBenefit = incurred
 263                       , _benefit = incurred + predicted
 264                       }
 265
 266 lastFlow :: Q.Seq SearchState -> Int -> CaveContext Int
 267 lastFlow Q.Empty _ = return 0
 268 lastFlow (_ :|> previous) newTime =
 269   do cf <- currentFlow previous
 270      let dt = newTime - (timeOf previous)
 271      return (cf * dt)
 272
 273 isGoal :: Agendum -> CaveContext Bool
 274 isGoal agendum =
 275   do timeLimit <- asks getTimeLimit
 276      let s = agendum ^. current
 277      return $ (timeOf s) == timeLimit
 278
 279 isFullFlow :: SearchState -> CaveContext Bool
 280 isFullFlow state =
 281   do cave <- asks getCave
 282      cf <- currentFlow state
 283      let ff = sumOf (folded . flowRate) cave
 284      return (cf == ff)
 285
 286 compress :: Cave -> Cave
 287 compress cave = M.mapWithKey (compressRoom cave) cave
 288
 289 compressRoom :: Cave -> RoomID -> Room -> Room
 290 compressRoom cave here room = room & tunnels .~ t'
 291   where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
 292
 293 reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
 294 reachableFrom _ [] _ routes = routes
 295 reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
 296   | here `S.member` found = reachableFrom cave boundary found routes
 297   | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
 298   where exits = (cave ! here) ^. tunnels
 299         exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
 300         legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
 301         routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
 302                   then routes
 303                   else S.insert tunnel routes
 304
 305 -- Parse the input file
 306
 307 caveP :: Parser Cave
 308 valveP :: Parser (RoomID, Room)
 309 roomP :: Parser Room
 310 tunnelsP :: Parser (S.Set Tunnel)
 311 tunnelTextP :: Parser Text
 312
 313 caveP = M.fromList <$> valveP `sepBy` endOfLine
 314 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
 315 roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
 316     -- where roomify v ts = Room {flowRate = v, tunnels = ts }
 317 tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
 318 tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
 319
 320 successfulParse :: Text -> Cave
 321 successfulParse input =
 322   case parseOnly caveP input of
 323     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 324     Right cave -> cave