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