2024 Day Three

DayThree.hs

{-# LANGUAGE OverloadedStrings #-}

module DayThree (partOne, partTwo, ParseState(Start), parseAllOps, sampleOps) where

import Safe (headDef)
import qualified Data.Maybe as M
import qualified Data.Text as T
import Data.Text (isPrefixOf)
import qualified Text.Read as TR
import Numeric (lexDigits)

data ParseState = Start | AwaitingFirstArg | AwaitingComma Int | AwaitingSecondArg Int | AwaitingCloseParen Int Int deriving (Show, Eq, Ord)
data Operation = MultiplicationOp Int Int | Activate | Deactivate deriving (Show, Eq)


sampleOps :: T.Text
sampleOps = "xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))"


parseOneOp :: ParseState -> T.Text -> (M.Maybe Operation, T.Text)
parseOneOp Start body
  | "do()" `isPrefixOf` body = (Just Activate, T.drop 4 body)
  | "don't()" `isPrefixOf` body = (Just Deactivate, T.drop 7 body)
  | "mul(" `isPrefixOf` body = parseOneOp AwaitingFirstArg (T.drop 4 body)
  | otherwise = (Nothing, T.drop 1 body)

parseOneOp AwaitingFirstArg body = go maybeDigits
  where go (Just n) = parseOneOp (AwaitingComma n) (T.drop (length (show n)) body)
        go _ = (Nothing, T.drop 1 body)
        maybeDigits = leadingDigitsString body

parseOneOp (AwaitingComma firstArg) body
  | "," `isPrefixOf` body = parseOneOp (AwaitingSecondArg firstArg) (T.drop 1 body)
  | otherwise = (Nothing, T.drop 1 body)

parseOneOp (AwaitingSecondArg firstArg) body = go maybeDigits
  where go (Just n) = parseOneOp (AwaitingCloseParen firstArg n) (T.drop (length (show n)) body)
        go _ = (Nothing, T.drop 1 body)
        maybeDigits = leadingDigitsString body


parseOneOp (AwaitingCloseParen firstArg secondArg)  body
  | ")" `isPrefixOf` body = (Just (MultiplicationOp firstArg secondArg), T.drop 1 body)
  | otherwise = (Nothing, T.drop 1 body)


parseAllOps :: T.Text -> [Operation]
parseAllOps "" = []
parseAllOps body = iterateParse firstFound
  where (firstFound, restBody) = parseOneOp Start body
        iterateParse Nothing = parseAllOps restBody
        iterateParse (Just op) = op:parseAllOps restBody


leadingDigitsString :: T.Text -> Maybe Int
leadingDigitsString body
  | not (null digitsStr) && length digitsStr <= 3 = Just (TR.read digitsStr)
  | otherwise = Nothing
  where (digitsStr, _) = headDef ("", "") (lexDigits (T.unpack body))

executeProgram :: [Operation] -> Int
executeProgram operations = executeRestOfProgram True operations
  where executeRestOfProgram _ [] = 0
        executeRestOfProgram _ (Activate:restOps) = executeRestOfProgram True restOps
        executeRestOfProgram _ (Deactivate:restOps) = executeRestOfProgram False restOps
        executeRestOfProgram True ((MultiplicationOp x y):restOps) = (x*y) + (executeRestOfProgram True restOps)
        executeRestOfProgram False ((MultiplicationOp _ _):restOps) = (executeRestOfProgram False restOps)

partOne :: T.Text -> Int
partOne body = sum $ map execute operations
  where
    operations = parseAllOps body
    execute (MultiplicationOp x y) = x * y
    execute _ = 0

partTwo :: T.Text -> Int
partTwo body = executeProgram $ parseAllOps body