AoC 2021, day 12

Day12.hs

module Day12 where

import Control.Applicative (Alternative((<|>)))
import Data.List ((\\))
import Data.Maybe (isJust)
import Data.Text (Text)
import GHC.Generics (Generic)
import Optics
import Text.Megaparsec (some)
import Text.Megaparsec.Char (char, eol, lowerChar, upperChar)

import qualified Data.Text as T
import qualified Data.Text.IO as TIO

import Parser ( parse, Parser )

data Cave = BigCave Text | SmallCave Text deriving (Eq, Ord, Show, Generic)

type Edge = (Cave, Cave)
type Path = [Cave]

part1 :: FilePath -> IO Int
part1 filename = solver filename nf
  where
    nf :: Path -> Bool -> [Cave] -> [Cave]
    nf path _ ns = ns \\ filter (isJust . preview #_SmallCave) path

part2 :: FilePath -> IO Int
part2 filename = solver filename nf
  where nf path flag ns | flag = ns \\ filter (isJust . preview #_SmallCave) path
                        | otherwise = ns \\ [SmallCave "start"]

solver :: FilePath -> (Path -> Bool -> [Cave] -> [Cave]) -> IO Int
solver filename nf = do
  edges <- parse (some parseEdge) =<< TIO.readFile filename
  pure . length $ findPaths edges (SmallCave "start") [] False nf

parseEdge :: Parser Edge
parseEdge = (,) <$> (cave <* char '-') <*> (cave <* eol)
  where cave = (SmallCave . T.pack <$> some lowerChar) <|> (BigCave . T.pack <$> some upperChar)

findPaths :: [Edge] -> Cave -> Path -> Bool -> (Path -> Bool -> [Cave] -> [Cave]) -> [Path]
findPaths edges cave path flag nf
  | cave == SmallCave "end" = [path']
  | otherwise = foldMap (\n -> findPaths edges n path' flag' nf) ns'
  where
    path' = cave : path
    flag' = flag || isJust (cave ^? #_SmallCave) && cave `elem` path
    ns'   = nf path flag' (neighbors edges cave)

neighbors :: Eq a => [(a, a)] -> a -> [a]
neighbors edges vertex = [ v2 | (v1, v2) <- edges, vertex == v1 ] ++ [ v1 | (v1, v2) <- edges, vertex == v2 ]