Mutable day 21

Day21.reddit.hs

#!/usr/bin/env stack
{-
    stack
    script
    --resolver lts-12.20
    --package text,trifecta,containers,parsers,mtl,vector,deepseq
-}
{-# LANGUAGE RecordWildCards, BangPatterns #-}
{-# OPTIONS_GHC -O2 #-}

import           Control.Applicative         ((<|>))
import           Data.Bits                   ((.&.), (.|.))
import qualified Data.IntSet                 as IntSet
import           Data.Vector                 (Vector)
import qualified Data.Vector                 as Vector
import qualified Data.Vector.Generic.Mutable as M
import           Data.Vector.Unboxed         ((!))
import qualified Data.Vector.Unboxed         as Unboxed
import           Debug.Trace
import           Text.Trifecta

import qualified Data.Vector.Unboxed.Mutable as MV
import           Data.IORef
import Control.Monad

type Registers = MV.IOVector Int

type Register = Int

data Instruction = Instruction
  { kind :: !OpKind
  , a    :: !Int
  , b    :: !Int
  , out  :: !Register
  } deriving (Show, Eq)

type IP = (Register, Int)

data OpKind
  = Addr
  | Addi
  | Mulr
  | Muli
  | Banr
  | Bani
  | Borr
  | Bori
  | Setr
  | Seti
  | Gtir
  | Gtri
  | Gtrr
  | Eqri
  | Eqir
  | Eqrr
  deriving (Show, Eq)

runProgram :: IP -> Vector Instruction -> Registers -> IO (Maybe Int)
runProgram (register, initialValue) instructions registers =
  go IntSet.empty initialValue 0
  where
    go :: IntSet.IntSet -> Int -> Int -> IO (Maybe Int)
    go seen pointerValue lastSolution = do
      MV.unsafeWrite registers register pointerValue
      Vector.unsafeIndex program pointerValue -- <- run the function

      nextPointerValue <- (+ 1) <$> MV.unsafeRead registers register :: IO Int
            -- ^ Retrieve register with pointer value, increment
      if pointerValue == 28
            then do
                value <- MV.unsafeRead registers 5
                if IntSet.member value seen
                    then return $ Just lastSolution
                    else    let seen' = IntSet.insert value seen
                            in  --trace (show $ IntSet.size seen') $
                                go
                                    seen'
                                    nextPointerValue
                                    value
            else go seen nextPointerValue lastSolution

    program = Vector.map run instructions :: Vector (IO ())
    rs = registers
    run :: Instruction -> IO ()
    run Instruction {..} = do
    --   traceM (show kind)
      x <- if (a < 0 || a >= M.basicLength rs) then return 0 else MV.unsafeRead rs a
      y <- if (b < 0 || b >= M.basicLength rs) then return 0 else MV.unsafeRead rs b
      let result =
            case kind of
              Addr -> x + y
              Addi -> x + b
              Mulr -> x * y
              Muli -> x * b
              Banr -> x .&. y
              Bani -> x .&. b
              Borr -> x .|. y
              Bori -> x .|. b
              Setr -> x
              Seti -> a
              Gtir ->
                if a > y
                  then 1
                  else 0
              Gtri ->
                if x > b
                  then 1
                  else 0
              Gtrr ->
                if x > y
                  then 1
                  else 0
              Eqri ->
                if x == b
                  then 1
                  else 0
              Eqir ->
                if a == y
                  then 1
                  else 0
              Eqrr ->
                if x == y
                  then 1
                  else 0
      -- in Unboxed.modify (\vec -> M.unsafeWrite vec out result) rs
      MV.unsafeWrite rs out result

main :: IO ()
main = do
  input <- parseString inputP mempty <$> readFile "data.txt"
  case input of
    Failure parseError -> print parseError
    Success (ip, instructions) -> do
      registers <- MV.new (length [0, 0, 0, 0, 0, 0])
      result <- runProgram ip instructions registers
      print result

-- Parsing stuff
instructionPointerP :: Parser IP
instructionPointerP =
  (,) <$> (string "#ip" *> whiteSpace *> (fromIntegral <$> natural)) <*> pure 0

instructionP :: Parser Instruction
instructionP = Instruction <$> opKindP <*> number' <*> number' <*> number'
  where
    number' = fromIntegral <$> (whiteSpace *> natural)
    opKindP =
      Addr <$ string "addr" <|> Addi <$ string "addi" <|> Mulr <$ string "mulr" <|>
      Muli <$ string "muli" <|>
      Banr <$ string "banr" <|>
      Bani <$ string "bani" <|>
      Borr <$ string "borr" <|>
      Bori <$ string "bori" <|>
      Setr <$ string "setr" <|>
      Seti <$ string "seti" <|>
      Gtir <$ string "gtir" <|>
      Gtri <$ string "gtri" <|>
      Gtrr <$ string "gtrr" <|>
      Eqir <$ string "eqir" <|>
      Eqri <$ string "eqri" <|>
      Eqrr <$ string "eqrr"

inputP :: Parser (IP, Vector Instruction)
inputP =
  (,) <$> instructionPointerP <* whiteSpace <*>
  (Vector.fromList <$> many instructionP)