advent19/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-19/
   2 -- Optimised at https://work.njae.me.uk/2023/07/24/optimising-haskell-example-4/
   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 qualified Data.MultiSet as MS
  17 import Data.Sequence ((|>))
  18 import Data.List
  19 import Data.Maybe
  20 -- import Data.Ord
  21 import Control.Monad.Reader
  22 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  23 import GHC.Generics (Generic)
  24 import Control.Parallel.Strategies
  25 import Control.DeepSeq
  26
  27
  28 data Resource = Ore | Clay | Obsidian | Geode
  29     deriving (Show, Eq, Ord, Generic, Enum, Bounded)
  30
  31 instance NFData Resource
  32
  33 type Collection = MS.MultiSet Resource
  34
  35 hashCollection :: Collection -> Int
  36 hashCollection c =
  37     let k = 200
  38     in sum $ zipWith (\r n -> (MS.occur r c) * n) [minBound .. maxBound] $ iterate (* k) 1
  39
  40 type Blueprint = M.Map Resource Collection
  41
  42 data TimedBlueprint = TimedBlueprint { getBlueprint :: Blueprint, getTimeLimit :: Int, getMaxRobots :: Collection}
  43     deriving (Show, Eq, Ord)
  44
  45 type BlueprintContext = Reader TimedBlueprint
  46
  47 data SearchState = SearchState
  48     { _resources :: Collection
  49     , _robots :: Collection
  50     , _currentTime :: Int
  51     } deriving (Eq, Show, Ord)
  52 makeLenses ''SearchState
  53
  54 type StateHash = (Int, Int, Int)
  55 hashSearchState :: SearchState -> StateHash
  56 hashSearchState s = (hashCollection (s ^. resources), hashCollection (s ^. robots), s ^. currentTime)
  57
  58 instance NFData SearchState where
  59   rnf (SearchState a b c) = rnf a `seq` rnf b `seq` rnf c `seq` ()
  60
  61 data Agendum =
  62     Agendum { _current :: SearchState
  63             , _trail :: Q.Seq SearchState
  64             , _trailBenefit :: Int
  65             , _benefit :: Int
  66             } deriving (Show, Eq, Ord)
  67 makeLenses ''Agendum
  68
  69 type Agenda = P.MaxPQueue Int Agendum
  70
  71 type ExploredStates = S.Set StateHash
  72
  73 main :: IO ()
  74 main =
  75   do  dataFileName <- getDataFileName
  76       text <- TIO.readFile dataFileName
  77       let blueprints = successfulParse text
  78       print $ part1 blueprints
  79       print $ part2 blueprints
  80
  81 part1, part2 :: [(Int, Blueprint)] -> Int
  82 part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- results]
  83 -- part1 blueprints = results
  84   -- where results = fmap (scoreBlueprint 24) blueprints
  85   where results = parMap rdeepseq (scoreBlueprint 24) blueprints
  86
  87 -- part2 :: [(Int, Blueprint)] -> Int
  88 part2 blueprints = product [MS.occur Geode (r ^. resources) | (_, r) <- results]
  89     -- where results = fmap (scoreBlueprint 32) $ take 3 blueprints
  90     where results = parMap rdeepseq (scoreBlueprint 32) $ take 3 blueprints
  91
  92 robotLimits :: Blueprint -> Collection
  93 robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
  94
  95 scoreBlueprint :: Int -> (Int, Blueprint) -> (Int, SearchState)
  96 scoreBlueprint t (n, bp) = ( n
  97                            , runReader searchSpace (TimedBlueprint bp t (robotLimits bp))
  98                            )
  99
 100 searchSpace :: BlueprintContext SearchState
 101 searchSpace =
 102     do agenda <- initAgenda
 103        -- searchAll agenda S.empty emptySearchState
 104        result <- aStar agenda S.empty
 105        return $ (fromJust result) ^. current
 106
 107 initAgenda :: BlueprintContext Agenda
 108 initAgenda =
 109     do let startState = emptySearchState
 110        b <- estimateBenefit startState
 111        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 112
 113
 114 aStar ::  Agenda -> ExploredStates -> BlueprintContext (Maybe Agendum)
 115 aStar agenda closed
 116     | P.null agenda = return Nothing
 117     | otherwise =
 118         do  let (_, currentAgendum) = P.findMax agenda
 119             let reached = currentAgendum ^. current
 120             nexts <- candidates currentAgendum closed
 121             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 122             let cl = hashSearchState reached
 123             atTimeLimit <- isTimeLimit currentAgendum
 124             if atTimeLimit
 125             then return (Just currentAgendum)
 126             else if (cl `S.member` closed)
 127                  then aStar (P.deleteMax agenda) closed
 128                  else aStar newAgenda (S.insert cl closed)
 129
 130 candidates :: Agendum -> ExploredStates -> BlueprintContext (Q.Seq Agendum)
 131 candidates agendum closed =
 132     do  let candidate = agendum ^. current
 133         let previous = agendum ^. trail
 134         succs <- successors candidate
 135         succAgs <- mapM (makeAgendum previous) succs
 136         let nonloops = Q.filter (\s -> (hashSearchState $ s ^. current) `S.notMember` closed) succAgs
 137         return nonloops
 138
 139 makeAgendum :: Q.Seq SearchState -> SearchState -> BlueprintContext Agendum
 140 makeAgendum previous newState =
 141     do predicted <- estimateBenefit newState
 142        let newTrail = previous |> newState
 143        let incurred = (MS.occur Geode (newState ^. resources))
 144        return Agendum { _current = newState
 145                       , _trail = newTrail
 146                       , _trailBenefit = incurred
 147                       , _benefit = incurred + predicted
 148                       }
 149
 150 isTimeLimit :: Agendum -> BlueprintContext Bool
 151 isTimeLimit agendum =
 152   do timeLimit <- asks getTimeLimit
 153      return $ (agendum ^. current . currentTime) >= timeLimit
 154
 155 emptySearchState :: SearchState
 156 emptySearchState = SearchState { _resources = MS.empty, _robots = MS.singleton Ore, _currentTime = 0 }
 157
 158 successors :: SearchState -> BlueprintContext (Q.Seq SearchState)
 159 successors state =
 160   do blueprint <- asks getBlueprint
 161      maxRobots <- asks getMaxRobots
 162      timeLimit <- asks getTimeLimit
 163
 164      let robotSuccessors = Q.fromList $ catMaybes $ M.elems $ M.mapWithKey (handleRobot state maxRobots timeLimit) blueprint
 165
 166      let timeRemaining = timeLimit - (state ^. currentTime)
 167      let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * timeRemaining) acc)
 168                                 MS.empty
 169                                 (state ^. robots)
 170      let delayUntilEnd = (state & currentTime .~ timeLimit
 171                                 & resources %~ (MS.union gathered)
 172                          )
 173      return ( robotSuccessors |> delayUntilEnd )
 174
 175 handleRobot :: SearchState -> Collection -> Int -> Resource -> Collection -> Maybe SearchState
 176 handleRobot state maxRobots timeLimit robot recipe
 177   | sufficientRobots robot state maxRobots = Nothing
 178   | otherwise = buildWhenReady robot state recipe timeLimit
 179
 180 -- do I already have enough of this robot?
 181 sufficientRobots :: Resource -> SearchState -> Collection -> Bool
 182 sufficientRobots robot state maxRobots =
 183     (robot `MS.member` maxRobots)
 184     &&
 185     ((MS.occur robot (state ^. robots)) >= (MS.occur robot maxRobots))
 186
 187 -- assuming I can't build this robot, how long do I have to wait for the current
 188 -- robots to gather enough to build it?
 189 buildDelay :: SearchState -> Collection -> Maybe Int
 190 buildDelay state recipe
 191   -- | MS.null delay = Just 0
 192   | all (\r -> MS.member r rbts) (MS.distinctElems shortfall) = Just $ maximum0 $ fmap snd $ MS.toOccurList delay
 193   | otherwise = Nothing
 194   where shortfall = recipe `MS.difference` (state ^. resources)
 195         delay = MS.foldOccur calcOneDelay MS.empty shortfall
 196         rbts = state ^. robots
 197         calcOneDelay resource count acc =
 198           MS.insertMany resource
 199                         -- (count `div` (MS.occur resource rbts) + 1)
 200                         (ceiling $ (fromIntegral count) / (fromIntegral $ MS.occur resource rbts))
 201                         acc
 202         maximum0 xs = if (null xs) then 0 else maximum xs
 203
 204 buildWhenReady :: Resource -> SearchState -> Collection -> Int -> Maybe SearchState
 205 buildWhenReady robot state recipe timeLimit =
 206   do  waitDelay <- buildDelay state recipe
 207       delay <- tooLate (state ^. currentTime) (waitDelay + 1) timeLimit
 208       let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * delay) acc)
 209                                 MS.empty
 210                                 (state ^. robots)
 211       return (state & robots %~ MS.insert robot                  -- add the robot
 212                     & resources %~ (MS.union gathered)           -- add the gathered resources
 213                     & resources %~ ( `MS.difference` recipe )    -- remove the resources to build it
 214                     & currentTime %~ (+ delay)
 215              )
 216
 217 tooLate :: Int -> Int -> Int -> Maybe Int
 218 tooLate current delay timeLimit
 219   | (current + delay) <= timeLimit = Just delay
 220   | otherwise = Nothing
 221
 222
 223 estimateBenefit :: SearchState -> BlueprintContext Int
 224 estimateBenefit currentState =
 225   do timeLimit <- asks getTimeLimit
 226      let timeElapsed = currentState ^. currentTime
 227      let timeRemaining = timeLimit - timeElapsed
 228      let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
 229      let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
 230      return $ currentRobotsGather + newRobotsGather
 231
 232
 233 -- Parse the input file
 234
 235 blueprintsP :: Parser [(Int, Blueprint)]
 236 blueprintP :: Parser (Int, Blueprint)
 237 robotP :: Parser (Resource, Collection)
 238 requirementsP :: Parser Collection
 239 requirementP :: Parser (Resource, Int)
 240 resourceP, oreP, clayP, obsidianP, geodeP :: Parser Resource
 241
 242 blueprintsP = blueprintP `sepBy` endOfLine
 243 blueprintP = blueprintify <$> (("Blueprint " *> decimal) <* ": ") <*> (robotP `sepBy` ". ") <* "."
 244   where blueprintify n robots =
 245           (n, M.fromList robots)
 246 robotP = (,) <$> ("Each " *> resourceP) <*> (" robot costs " *> requirementsP)
 247
 248 requirementsP = MS.fromOccurList <$> (requirementP `sepBy` " and ")
 249
 250 requirementP = (flip (,)) <$> (decimal <* " ") <*> resourceP
 251
 252 resourceP = oreP <|> clayP <|> obsidianP <|> geodeP
 253 oreP = Ore <$ "ore"
 254 clayP = Clay <$ "clay"
 255 obsidianP = Obsidian <$ "obsidian"
 256 geodeP = Geode <$ "geode"
 257
 258 successfulParse :: Text -> [(Int, Blueprint)]
 259 successfulParse input =
 260   case parseOnly blueprintsP input of
 261     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 262     Right blueprints -> blueprints