{-# Language BlockArguments, NumericUnderscores, QuasiQuotes #-}
{-|
Module      : Main
Description : Day 15 solution
Copyright   : (c) Eric Mertens, 2020
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2020/day/15>

-}
module Main (main) where

import Advent.Format (format)
import Control.Monad (zipWithM_)
import Control.Monad.ST (ST, runST)
import Data.Int (Int32)
import Data.Primitive.PrimArray (MutablePrimArray, readPrimArray, writePrimArray, newPinnedPrimArray, setPrimArray)

-- | Type of elements in our sequence -- big enough to hold 30 million
type T = Int32

-- |
-- >>> game [10,16,6,0,1,17] 2020
-- 412
main :: IO ()
IO ()
main =
  do inp <- (Int -> T) -> [Int] -> [T]
forall a b. (a -> b) -> [a] -> [b]
map Int -> T
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> [T]) -> IO [Int] -> IO [T]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2020 15 %u&,%n|]
     print (game inp      2_020)
     print (game inp 30_000_000)

game ::
  [T] {- ^ initial sequence -} ->
  T   {- ^ desired position -} ->
  T   {- ^ desired element  -}
game :: [T] -> T -> T
game [T]
xs T
n = (forall s. ST s T) -> T
forall a. (forall s. ST s a) -> a
runST
  do let len :: Int
len = T -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (T -> T -> T
forall a. Ord a => a -> a -> a
max T
n (T
1 T -> T -> T
forall a. Num a => a -> a -> a
+ [T] -> T
forall a. HasCallStack => [a] -> a
last [T]
xs))
     a <- Int -> ST s (MutablePrimArray (PrimState (ST s)) T)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPinnedPrimArray Int
len
     setPrimArray a 0 len 0
     zipWithM_ (writePrimArray a) (fromIntegral <$> Prelude.init xs) [1..]
     speak a n (fromIntegral (length xs)) (last xs)

speak ::
  MutablePrimArray s T {- ^ position of last occurrence -} ->
  T      {- ^ desired position -} ->
  T      {- ^ current position -} ->
  T      {- ^ current element  -} ->
  ST s T {- ^ desired element  -}
speak :: forall s. MutablePrimArray s T -> T -> T -> T -> ST s T
speak MutablePrimArray s T
a T
n T
m T
x
  | T
m T -> T -> Bool
forall a. Eq a => a -> a -> Bool
== T
n    = T -> ST s T
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (T -> ST s T) -> T -> ST s T
forall a b. (a -> b) -> a -> b
$! T
x
  | Bool
otherwise = do v <- MutablePrimArray s T -> Int -> T -> ST s T
forall s. MutablePrimArray s T -> Int -> T -> ST s T
exchange MutablePrimArray s T
a (T -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral T
x) T
m
                   speak a n (m+1) (if v == 0 then 0 else m-v)

-- | Exchange element at an index with a new element returning old element.
exchange :: MutablePrimArray s T -> Int -> T -> ST s T
exchange :: forall s. MutablePrimArray s T -> Int -> T -> ST s T
exchange MutablePrimArray s T
a Int
i T
x = MutablePrimArray (PrimState (ST s)) T -> Int -> ST s T
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray s T
MutablePrimArray (PrimState (ST s)) T
a Int
i ST s T -> ST s () -> ST s T
forall a b. ST s a -> ST s b -> ST s a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* MutablePrimArray (PrimState (ST s)) T -> Int -> T -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s T
MutablePrimArray (PrimState (ST s)) T
a Int
i T
x