0 | module Days.Day4

  2 | import Data.String

  3 | import Data.Vect

  4 | import Data.Fin

  5 | import Data.List.Lazy

  6 | import Control.Monad.State

  7 | import System

  9 | record GridCell where

 10 |   constructor MkGridCell

 11 |   content : Char

 12 |   visit_count : Nat

 14 | fromChar : Char -> GridCell

 15 | fromChar c = MkGridCell c 0

 17 | data Grid : Nat -> Nat -> Type where

 18 |   MkGrid : {rows, cols : Nat} -> Vect (S rows) (Vect (S cols) GridCell)

 19 |          -> Grid (S rows) (S cols)

 21 | %name Grid grid, grid2, grid3

 23 | index : Grid (S rows) (S cols) -> (Fin (S rows), Fin (S cols)) -> GridCell

 24 | index (MkGrid xs) (row, col) =

 25 |   let row' = row `index` xs

 26 |       out = col `index` row'

 27 |   in out

 29 | mark : Grid (S rows) (S cols) -> (Fin (S rows), Fin (S cols)) -> Grid (S rows) (S cols)

 30 | mark (MkGrid xs) (row, col) =

 31 |   let row' = row `index` xs

 32 |       orig = col `index` row'

 33 |       new : GridCell = { visit_count $= S } orig

 34 |       row'' = replaceAt col new row'

 35 |   in MkGrid $ replaceAt row row'' xs

 37 | coords : {rows, cols : Nat} -> LazyList (Fin (S rows), Fin (S cols))

 38 | coords =

 39 |   let row_coords : List (Fin (S rows)) = allFins _

 40 |       col_coords : List (Fin (S cols)) = allFins _

 41 |   in concat . map (mkRow col_coords) $ row_coords

 42 |   where

 43 |     mkRow : (cols_list : List (Fin (S cols))) -> (row : Fin (S rows))

 44 |           -> LazyList (Fin (S rows), Fin (S cols))

 45 |     mkRow [] row = []

 46 |     mkRow (x :: xs) row = (row, x) :: mkRow xs row

 48 | data Dir = Inc | Dec | Nop

 50 | directions : List (Dir, Dir)

 51 | directions = [(Inc, Inc), (Inc, Dec), (Inc, Nop),

 52 |               (Dec, Inc), (Dec, Dec), (Dec, Nop),

 53 |               (Nop, Inc), (Nop, Dec)]

 55 | applyDir : {x : Nat} -> (dir : Dir) -> (coord : Fin (S x)) -> Maybe (Fin (S x))

 56 | applyDir Inc coord =

 57 |   if coord < last

 58 |   then Just $ finS coord

 59 |   else Nothing

 60 | applyDir Dec FZ = Nothing

 61 | applyDir Dec (FS y) = Just . weaken $ y

 62 | applyDir Nop coord = Just coord

 64 | applyDirs : {rows, cols : Nat} -> (dir_pair : (Dir, Dir))

 65 |           -> (coord_pair : (Fin (S rows), Fin (S cols)))

 66 |           -> Maybe (Fin (S rows), Fin (S cols))

 67 | applyDirs (row_dir, col_dir) (row_coord, col_coord) = do

 68 |   row_coord <- applyDir row_dir row_coord

 69 |   col_coord <- applyDir col_dir col_coord

 70 |   Just (row_coord, col_coord)

 72 | Coord : (rows, cols : Nat) -> Type

 73 | Coord rows cols = (Fin (S rows), Fin (S cols))

 75 | Ray : (rows, cols, len : Nat) -> Type

 76 | Ray rows cols len = Vect len (Coord rows cols)

 78 | ray : {rows, cols : Nat} -> (dir_pair : (Dir, Dir)) -> (coord_pair : (Coord rows cols))

 79 |     -> (count : Nat) -> Maybe (Ray rows cols count)

 80 | ray dir_pair coord_pair 0 = Just []

 81 | ray dir_pair coord_pair (S 0) = Just [coord_pair]

 82 | ray dir_pair coord_pair (S (S k)) = do

 83 |   next_coord <- applyDirs dir_pair coord_pair

 84 |   rest <- ray dir_pair next_coord (S k)

 85 |   Just $ coord_pair :: rest

 87 | rays : {rows, cols : Nat} -> (coord_pair : (Coord rows cols)) -> (count : Nat)

 88 |      -> List (Ray rows cols count)

 89 | rays coord_pair count = mapMaybe (\x => ray x coord_pair count) directions

 91 | indexRay : {rows, cols : Nat} -> (ray : Ray rows cols len) -> Grid (S rows) (S cols)

 92 |          -> Vect len GridCell

 93 | indexRay [] grid = []

 94 | indexRay (x :: xs) grid =

 95 |  (grid `index` x) :: indexRay xs grid

 97 | data XDir = IncInc | DecInc | IncDec | DecDec

 99 | xDirs : List (XDir, XDir)

100 | xDirs = [(IncInc, IncDec), (IncInc, DecInc), (DecDec, IncDec), (DecDec, DecInc)]

102 | finDec : Fin (S x) -> Fin (S x)

103 | finDec FZ = FZ

104 | finDec (FS y) = weaken y

106 | xray : {rows, cols : Nat} -> (dir: XDir) -> (coord_pair : (Coord rows cols))

107 |      -> Maybe (Ray rows cols 3)

108 | xray dir (row, col) =

109 |   if row > 0 && col > 0 && row < last && col < last

110 |   then case dir of

111 |     IncInc => Just [(finDec row, finDec col), (row, col), (finS row, finS col)]

112 |     DecInc => Just [(finS row, finDec col), (row, col), (finDec row, finS col)]

113 |     IncDec => Just [(finDec row, finS col), (row, col), (finS row, finDec col)]

114 |     DecDec => Just [(finS row, finS col), (row, col), (finDec row, finDec col)]

115 |   else Nothing

117 | xrayPair : {rows, cols : Nat} -> (coord_pair : Coord rows cols) -> (dirs : (XDir, XDir))

118 |          -> Maybe (Ray rows cols 3, Ray rows cols 3)

119 | xrayPair coord_pair (x, y) = do

120 |   x' <- xray x coord_pair

121 |   y' <- xray y coord_pair

122 |   Just (x', y')

124 | xrayPairs : {rows, cols : Nat} -> (coord_pair : Coord rows cols)

125 |           -> List (Ray rows cols 3, Ray rows cols 3)

126 | xrayPairs coord_pair = mapMaybe (xrayPair coord_pair) xDirs

128 | rayToString : Grid (S rows) (S cols) -> Ray rows cols len -> String

129 | rayToString grid xs = pack . toList . map (content . index grid) $ xs

131 | markRay : Grid (S rows) (S cols) -> Ray rows cols len

132 |         -> Grid (S rows) (S cols)

133 | markRay grid [] = grid

134 | markRay grid (x :: xs) = markRay (mark grid x) xs

136 | checkRay : {rows, cols : Nat} -> (needle : String) -> (len : Nat)

137 |          -> (ray : Ray rows cols len)

138 |          -> State (Grid (S rows) (S cols)) ()

139 | checkRay needle len ray = do

140 |   grid <- get

141 |   let rs = rayToString grid ray

142 |   if rs == needle

143 |     then do

144 |       let grid = markRay grid ray

145 |       put grid

146 |     else pure ()

148 | checkCoord : {rows, cols : Nat} -> (needle : String) -> (coords : Coord rows cols)

149 |            -> State (Grid (S rows) (S cols)) ()

150 | checkCoord needle coords = do

151 |   grid <- get

152 |   let len = length . unpack $ needle

153 |   traverse_ (checkRay needle len) (rays coords len)

155 | checkCoords : {rows, cols : Nat} -> (needle : String)

156 |             -> (coords : LazyList (Coord rows cols))

157 |             -> State (Grid (S rows) (S cols)) ()

158 | checkCoords needle [] = pure ()

159 | checkCoords needle (x :: xs) = do

160 |   checkCoord needle x

161 |   checkCoords needle xs

163 | markX : Grid (S rows) (S cols) -> (Ray rows cols 3, Ray rows cols 3) -> Grid (S rows) (S cols)

164 | markX grid (x, y) =

165 |   let (len ** ray) = nub $ x ++ y

166 |   in markRay grid ray

168 | checkX : {rows, cols : Nat} -> (needle : Vect 3 Char) -> (Ray rows cols 3, Ray rows cols 3)

169 |        -> State (Grid (S rows) (S cols)) ()

170 | checkX needle (x, y) = do

171 |   grid <- get

172 |   let x' = map content . indexRay x $ grid

173 |   let y' = map content . indexRay y $ grid

174 |   if x' == needle && y' == needle

175 |    then do

176 |      let grid = markX grid (x, y)

177 |      put grid

178 |    else pure ()

180 | checkCoordX : {rows, cols : Nat} -> (needle : Vect 3 Char) -> (coords : Coord rows cols)

181 |             -> State (Grid (S rows) (S cols)) ()

182 | checkCoordX needle coords = do

183 |   grid <- get

184 |   let pairs = xrayPairs

185 |   traverse_ (checkX needle) (xrayPairs coords)

187 | checkCoordsX : {rows, cols : Nat} -> (needle : Vect 3 Char)

188 |              -> (coords : LazyList (Coord rows cols))

189 |              -> State (Grid (S rows) (S cols)) ()

190 | checkCoordsX needle [] = pure ()

191 | checkCoordsX needle (x :: xs) = do

192 |   checkCoordX needle x

193 |   checkCoordsX needle xs

195 | stringToGrid : String -> Maybe (rows ** cols ** Grid (S rows) (S cols))

196 | stringToGrid str = do

197 |   let cells = map (map fromChar . unpack) . map trim . lines . trim $ str

198 |   row <- head' cells

199 |   (S cols) <- Just $ length row

200 |     | _ => Nothing

201 |   cells <- traverse (toVect (S cols)) cells

202 |   (S rows) <- Just $ length cells

203 |     | _ => Nothing

204 |   cells <- toVect (S rows) cells

205 |   Just $ (rows ** cols ** MkGrid cells)

207 | showGrid : Grid (S rows) (S cols) -> String

208 | showGrid (MkGrid xs) = unlines . toList . map (pack . toList . map content) $ xs

210 | showCellHit : GridCell -> Char

211 | showCellHit (MkGridCell content visit_count) =

212 |   if visit_count > 0

213 |   then content

214 |   else '.'

216 | showHits : Grid (S rows) (S cols) -> String

217 | showHits (MkGrid xs) = unlines . toList . map (pack . toList . map showCellHit) $ xs

219 | countStarts : Grid (S rows) (S cols) -> (start_char : Char) -> Nat

220 | countStarts (MkGrid xs) start_char =

221 |     sum . map visit_count .

222 |     filter ((== start_char) . content) . toList . Data.Vect.concat $ xs

225 | part1 : String -> IO (String, ())

226 | part1 str =

227 |   case stringToGrid str of

228 |     Nothing => do

229 |       putStrLn "Failed to parse grid"

230 |       exitFailure

231 |     Just (rows ** cols ** grid) => do

232 |       let comp = checkCoords "XMAS" coords

233 |       let (grid, _) = runState grid comp

234 |       pure (show $ countStarts grid 'X', ())

237 | part2 : () -> String -> IO String

238 | part2 x str =

239 |   case stringToGrid str of

240 |     Nothing => do

241 |       putStrLn "Failed to parse grid"

242 |       exitFailure

243 |     Just (rows ** cols ** grid) => do

244 |       let comp = checkCoordsX ['M', 'A', 'S'] coords

245 |       let (grid, _) = runState grid comp

246 |       pure . show $ countStarts grid 'A'