{-# Language ImportQualifiedPost, BangPatterns, LambdaCase #-}
{-# OPTIONS_GHC -funbox-strict-fields #-} -- makes Photon more efficient
{-|
Module      : Main
Description : Day 16 solution
Copyright   : (c) Eric Mertens, 2023
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2023/day/16>

This is a pretty straight forward graph traversal of the
state space. I represent nodes of the "graph" being searched
as pairs of a location and a direction vector. At each step
the location is used to look up the tile and the direction
vector is used to compute reflections and splits.

Optimizations:

- Parallelize the search in part 2
- Only track seen states for beam splitters
- Count visited locations with an array instead of a Set
- Pack photon states into an Int to make seen set lookups faster

>>> :{
:main +
".|...\\....
|.-.\\.....
.....|-...
........|.
..........
.........\\
..../.\\\\..
.-.-/..|..
.|....-|.\\
..//.|....
"
:}
46
51

-}
module Main (main) where

import Advent (getInputArray, countBy)
import Advent.Coord (east, invert, invert', north, origin, south, west, coordCol, coordRow, Coord(C))
import Control.Parallel.Strategies (parMap, rpar)
import Data.Array.Unboxed (inRange, bounds, UArray, (!), accumArray, elems)
import Data.IntSet qualified as IntSet

-- | A single photon's location and velocity.
data Photon = P !Coord !Coord -- ^ location velocity

-- | Parse the input grid and print answers to both parts.
--
-- >>> :main
-- 7979
-- 8437
main :: IO ()
IO ()
main =
 do input <- Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2023 Int
16
    print (solve input (P origin east))
    print (maximum (parMap rpar (solve input) (edges (bounds input))))

-- | Count the number of energized tiles given an input beam.
solve :: UArray Coord Char -> Photon -> Int
solve :: UArray Coord Char -> Photon -> Int
solve UArray Coord Char
input = (Coord, Coord) -> [Photon] -> Int
coverage (UArray Coord Char -> (Coord, Coord)
forall i. Ix i => UArray i Char -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Char
input) ([Photon] -> Int) -> (Photon -> [Photon]) -> Photon -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Photon -> Bool)
-> (Photon -> Int) -> (Photon -> [Photon]) -> Photon -> [Photon]
forall a. (a -> Bool) -> (a -> Int) -> (a -> [a]) -> a -> [a]
dfsOn Photon -> Bool
isSplitter Photon -> Int
photonRep (UArray Coord Char -> Photon -> [Photon]
step UArray Coord Char
input)
  where
    -- branching only happens at splitters, so only bother avoiding
    -- duplication of work when visiting them
    isSplitter :: Photon -> Bool
isSplitter (P Coord
here Coord
_) =
      case UArray Coord Char
input UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
here of
        Char
'-' -> Bool
True
        Char
'|' -> Bool
True
        Char
_   -> Bool
False

-- | Use a more compact representative of the state space to speed
-- up the visited test. This saves about a 3rd of the runtime as without.
photonRep :: Photon -> Int
photonRep :: Photon -> Int
photonRep (P (C Int
y Int
x) (C Int
dy Int
dx)) = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4096 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
dyInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
dxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

-- | Find all the incoming light possibilities for part 2
edges :: (Coord, Coord) {- ^ bounds -} -> [Photon]
edges :: (Coord, Coord) -> [Photon]
edges (C Int
y1 Int
x1, C Int
y2 Int
x2) =
  [Coord -> Coord -> Photon
P (Int -> Int -> Coord
C Int
y1 Int
x) Coord
south | Int
x <- [Int
x1..Int
x2]] [Photon] -> [Photon] -> [Photon]
forall a. [a] -> [a] -> [a]
++
  [Coord -> Coord -> Photon
P (Int -> Int -> Coord
C Int
y2 Int
x) Coord
north | Int
x <- [Int
x1..Int
x2]] [Photon] -> [Photon] -> [Photon]
forall a. [a] -> [a] -> [a]
++
  [Coord -> Coord -> Photon
P (Int -> Int -> Coord
C Int
y Int
x1) Coord
east  | Int
y <- [Int
y1..Int
y2]] [Photon] -> [Photon] -> [Photon]
forall a. [a] -> [a] -> [a]
++
  [Coord -> Coord -> Photon
P (Int -> Int -> Coord
C Int
y Int
x2) Coord
west  | Int
y <- [Int
y1..Int
y2]]

-- | Advance a photon once cell forward and track its
-- resulting outgoing photons.
step :: UArray Coord Char -> Photon -> [Photon]
step :: UArray Coord Char -> Photon -> [Photon]
step UArray Coord Char
input (P Coord
here Coord
dir) =
  [ Coord -> Coord -> Photon
P Coord
here' Coord
dir'
  | Coord
dir' <-
      case UArray Coord Char
input UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
here of
        Char
'\\'                    -> [Coord -> Coord
invert Coord
dir]
        Char
'/'                     -> [Coord -> Coord
invert' Coord
dir]
        Char
'|' | Coord -> Int
coordRow Coord
dir Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> [Coord
north, Coord
south]
        Char
'-' | Coord -> Int
coordCol Coord
dir Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> [Coord
east, Coord
west]
        Char
_                       -> [Coord
dir]
  , let here' :: Coord
here' = Coord
here Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord
dir'
  , (Coord, Coord) -> Coord -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (UArray Coord Char -> (Coord, Coord)
forall i. Ix i => UArray i Char -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Char
input) Coord
here'
  ]

-- This is a copy of Advent.Search.dfsOn but augmented with
-- a predicate identifying the states that cause branching
-- so we don't bother deduplicating states that don't matter.
dfsOn :: (a -> Bool) -> (a -> Int) -> (a -> [a]) -> a -> [a]
dfsOn :: forall a. (a -> Bool) -> (a -> Int) -> (a -> [a]) -> a -> [a]
dfsOn a -> Bool
p a -> Int
rep a -> [a]
next a
start = IntSet -> [a] -> [a]
loop IntSet
IntSet.empty [a
start]
  where
    loop :: IntSet -> [a] -> [a]
loop !IntSet
seen = \case
      [] -> []
      a
x : [a]
q
        | Bool
slow, Int -> IntSet -> Bool
IntSet.member Int
r IntSet
seen ->     IntSet -> [a] -> [a]
loop IntSet
seen  [a]
q
        | Bool
otherwise                  -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: IntSet -> [a] -> [a]
loop IntSet
seen' [a]
q'
        where
          slow :: Bool
slow  = a -> Bool
p a
x
          r :: Int
r     = a -> Int
rep a
x
          seen' :: IntSet
seen' = if Bool
slow then Int -> IntSet -> IntSet
IntSet.insert Int
r IntSet
seen else IntSet
seen
          q' :: [a]
q'    = a -> [a]
next a
x [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
q

-- A more efficient way to count the number of unique coordinates
-- in the energized laser path.
coverage :: (Coord, Coord) {- ^ bounds -} -> [Photon] -> Int
coverage :: (Coord, Coord) -> [Photon] -> Int
coverage (Coord, Coord)
bnds [Photon]
path = (Bool -> Bool) -> [Bool] -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy Bool -> Bool
forall a. a -> a
id (UArray Coord Bool -> [Bool]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems UArray Coord Bool
a)
  where
    a :: UArray Coord Bool
    a :: UArray Coord Bool
a = (Bool -> () -> Bool)
-> Bool -> (Coord, Coord) -> [(Coord, ())] -> UArray Coord Bool
forall (a :: * -> * -> *) e i e'.
(IArray a e, Ix i) =>
(e -> e' -> e) -> e -> (i, i) -> [(i, e')] -> a i e
accumArray (\Bool
_ ()
_ -> Bool
True) Bool
False (Coord, Coord)
bnds [(Coord
p, ()) | P Coord
p Coord
_v <- [Photon]
path]