{-# 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 [T]
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|]
     T -> IO ()
forall a. Show a => a -> IO ()
print ([T] -> T -> T
game [T]
inp      T
2_020)
     T -> IO ()
forall a. Show a => a -> IO ()
print ([T] -> T -> T
game [T]
inp T
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))
     MutablePrimArray s T
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
     MutablePrimArray (PrimState (ST s)) T -> Int -> Int -> T -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray MutablePrimArray s T
MutablePrimArray (PrimState (ST s)) T
a Int
0 Int
len T
0
     (Int -> T -> ST s ()) -> [Int] -> [T] -> ST s ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (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) (T -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (T -> Int) -> [T] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [T] -> [T]
forall a. HasCallStack => [a] -> [a]
Prelude.init [T]
xs) [T
1..]
     MutablePrimArray s T -> T -> T -> T -> ST s T
forall s. MutablePrimArray s T -> T -> T -> T -> ST s T
speak MutablePrimArray s T
a T
n (Int -> T
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([T] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [T]
xs)) ([T] -> T
forall a. HasCallStack => [a] -> a
last [T]
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 T
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
                   MutablePrimArray s T -> T -> T -> T -> ST s T
forall s. MutablePrimArray s T -> T -> T -> T -> ST s T
speak MutablePrimArray s T
a T
n (T
mT -> T -> T
forall a. Num a => a -> a -> a
+T
1) (if T
v T -> T -> Bool
forall a. Eq a => a -> a -> Bool
== T
0 then T
0 else T
mT -> T -> T
forall a. Num a => a -> a -> a
-T
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