{-# Language QuasiQuotes, ImportQualifiedPost, NumDecimals #-}
{-|
Module      : Main
Description : Day 17 solution
Copyright   : (c) Eric Mertens, 2017
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2017/day/17>

Day 17 has us repeatedly insert elements into a circular buffer and
asks about the elements following some needle.

Part 1 is small enough that we can generate the whole list quickly
and search it.

Part 2 is large enough that generating the whole list and searching
it uses roughly 5.2 GB of RAM and takes 1 minute 15 seconds on my
hardware to run! We instead optimize things by noticing that the @0@
element is always at the head of the sequence, so we simply need to
find the last element that was written at index @1@.

-}
module Main where

import Advent (format)
import Data.List (elemIndices, foldl', scanl')
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq

-- | Print the solutions to the puzzle. Input file can be overridden
-- via command-line arguments.
--
-- >>> :main
-- 866
-- 11995607
main :: IO ()
IO ()
main =
  do Int
input <- [format|2017 17 %u%n|]
     Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> Seq Int -> Int
elemAfter Int
2017 (Int -> Int -> Seq Int
makeSequence Int
input Int
2017))
     Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> Int
part2 Int
input)

-- | Compute the element that immediately follows the needle
-- in the haystack. This assumes a circular interpretation for
-- the list, so the first element is considered to follow the last.
elemAfter ::
  Int     {- ^ needle            -} ->
  Seq Int {- ^ haystack          -} ->
  Int     {- ^ following element -}
elemAfter :: Int -> Seq Int -> Int
elemAfter Int
x Seq Int
xs = Seq Int -> Int -> Int
forall a. Seq a -> Int -> a
Seq.index Seq Int
xs ( (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Seq Int -> Int
forall a. Seq a -> Int
Seq.length Seq Int
xs )
  where
    Just Int
i = Int -> Seq Int -> Maybe Int
forall a. Eq a => a -> Seq a -> Maybe Int
Seq.elemIndexL Int
x Seq Int
xs

-- | Compute the buffer generated by inserting elements up to a given
-- element using a particular jump size.
--
-- >>> makeSequence 3 9
-- fromList [0,9,5,7,2,4,3,8,6,1]
makeSequence ::
  Int {- ^ jump size    -} ->
  Int {- ^ last element -} ->
  Seq Int
makeSequence :: Int -> Int -> Seq Int
makeSequence Int
jump Int
sz
  = (Seq Int -> (Int, Int) -> Seq Int)
-> Seq Int -> [(Int, Int)] -> Seq Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Seq Int
xs (Int
x,Int
i) -> Int -> Int -> Seq Int -> Seq Int
forall a. Int -> a -> Seq a -> Seq a
Seq.insertAt Int
i Int
x Seq Int
xs) Seq Int
forall a. Seq a
Seq.empty
  ([(Int, Int)] -> Seq Int) -> [(Int, Int)] -> Seq Int
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..Int
sz]
  ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
cursors Int
jump

-- | The infinite list of cursors generated from a particular jump parameter.
--
-- >>> take 10 (cursors 3)
-- [0,1,1,2,2,1,5,2,6,1]
cursors ::
  Int   {- ^ jump size        -} ->
  [Int] {- ^ cursor positions -}
cursors :: Int -> [Int]
cursors Int
jump = (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' Int -> Int -> Int
nextCursor Int
0 [Int
1..]
  where
    nextCursor :: Int -> Int -> Int
nextCursor Int
cursor Int
size = (Int
cursorInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
jump)Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem`Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
{-# Inline cursors #-} -- helps list fusion!

-- | Special case for when we only need to know what number is going
-- to follow the zero. Because the 0 is always going to be at the zero
-- index, whatever the last element to be written to the 1 index must
-- be the element that directly follows the zero.
part2 ::
  Int {- ^ jump size             -} ->
  Int {- ^ number following zero -}
part2 :: Int -> Int
part2 = [Int] -> Int
forall a. HasCallStack => [a] -> a
last ([Int] -> Int) -> (Int -> [Int]) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Eq a => a -> [a] -> [Int]
elemIndices Int
1 ([Int] -> [Int]) -> (Int -> [Int]) -> Int -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
5e7 ([Int] -> [Int]) -> (Int -> [Int]) -> Int -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int]
cursors