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

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

-}
module Main where

import Advent (format)
import Control.Monad.ST (ST, runST)
import Data.Vector.Unboxed qualified as V
import Data.Vector.Unboxed.Mutable qualified as M

main :: IO ()
IO ()
main =
  do [Int]
input <- [format|2017 5 (%d%n)*|]
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Int -> Int) -> [Int] -> Int
solve Int -> Int
part1 [Int]
input)
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Int -> Int) -> [Int] -> Int
solve Int -> Int
part2 [Int]
input)

-- | Update rules
part1, part2 :: Int -> Int
part1 :: Int -> Int
part1 Int
x             = Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
part2 :: Int -> Int
part2 Int
x | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3    = Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
        | Bool
otherwise = Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1

-- | Compute the number of steps until the program terminates given
-- an update rule.
--
-- >>> solve part1 [0,3,0,1,-3]
-- 5
-- >>> solve part2 [0,3,0,1,-3]
-- 10
solve ::
  (Int -> Int) {- ^ update rule     -} ->
  [Int]        {- ^ initial program -} ->
  Int          {- ^ steps required  -}
solve :: (Int -> Int) -> [Int] -> Int
solve Int -> Int
f [Int]
xs = (forall s. ST s Int) -> Int
forall a. (forall s. ST s a) -> a
runST (Int -> Int -> MVector (PrimState (ST s)) Int -> ST s Int
forall {f :: * -> *} {t}.
(PrimMonad f, Num t) =>
t -> Int -> MVector (PrimState f) Int -> f t
loop Int
0 Int
0 (MVector s Int -> ST s Int) -> ST s (MVector s Int) -> ST s Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
V.thaw ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
V.fromList [Int]
xs))
  where
    loop :: t -> Int -> MVector (PrimState f) Int -> f t
loop t
steps Int
i MVector (PrimState f) Int
mem
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= MVector (PrimState f) Int -> Int
forall a s. Unbox a => MVector s a -> Int
M.length MVector (PrimState f) Int
mem = t -> f t
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (t -> f t) -> t -> f t
forall a b. (a -> b) -> a -> b
$! t
steps
      | Bool
otherwise =
          do Int
d <- MVector (PrimState f) Int -> Int -> f Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
M.read MVector (PrimState f) Int
mem Int
i
             MVector (PrimState f) Int -> Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
M.write MVector (PrimState f) Int
mem Int
i (Int -> Int
f Int
d)
             t -> Int -> MVector (PrimState f) Int -> f t
loop (t
stepst -> t -> t
forall a. Num a => a -> a -> a
+t
1) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) MVector (PrimState f) Int
mem