{-|
Module:         Day16
Description:    <https://adventofcode.com/2021/day/16 Day 16: Packet Decoder>
-}
{-# LANGUAGE FlexibleContexts, LambdaCase, NamedFieldPuns, RecordWildCards, TypeApplications, TypeFamilies, ViewPatterns #-}
module Day16 (day16a, day16b) where

import Control.Monad.State (MonadState, evalStateT, gets, put, state)
import Data.Bits (testBit)
import Data.Char (digitToInt)
import Data.List (foldl')
import Data.Text (Text)
import Data.Void (Void)
import Text.Megaparsec (MonadParsec, ParseErrorBundle, Token, count, getOffset, parse)
import Text.Megaparsec.Char (hexDigitChar)

data Packet a
  = Literal { Packet a -> Int
version :: Int, Packet a -> a
value :: a }
  | Operator { version :: Int, Packet a -> Int
tag :: Int, Packet a -> [Packet a]
packets :: [Packet a] }

parser :: (Num a, MonadParsec e s m, Token s ~ Char) => m (Packet a)
parser :: m (Packet a)
parser = StateT [Bool] m (Packet a) -> [Bool] -> m (Packet a)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT [Bool] m (Packet a)
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m (Packet a)
packet [] where
    packet :: (Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) => m (Packet a)
    packet :: m (Packet a)
packet = do
        Int
version <- Int -> m Int
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m a
bits Int
3
        Int
tag <- Int -> m Int
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m a
bits Int
3
        if Int
tag Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4 then Int -> a -> Packet a
forall a. Int -> a -> Packet a
Literal Int
version (a -> Packet a) -> m a -> m (Packet a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m a
literal else do
            Bool
b <- m Bool
forall (m :: * -> *) e s.
(MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m Bool
bit
            [Packet a]
packets <- if Bool
b then Int -> m Int
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m a
bits Int
11 m Int -> (Int -> m [Packet a]) -> m [Packet a]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> m (Packet a) -> m [Packet a])
-> m (Packet a) -> Int -> m [Packet a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> m (Packet a) -> m [Packet a]
forall (m :: * -> *) a. Monad m => Int -> m a -> m [a]
count m (Packet a)
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m (Packet a)
packet else Int -> m Int
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m a
bits Int
15 m Int -> (Int -> m [Packet a]) -> m [Packet a]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> m [Packet a]
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m [Packet a]
parseSpan
            Packet a -> m (Packet a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Operator :: forall a. Int -> Int -> [Packet a] -> Packet a
Operator { Int
[Packet a]
packets :: [Packet a]
tag :: Int
version :: Int
packets :: [Packet a]
tag :: Int
version :: Int
.. }
    getHead :: (MonadState [a] m) => m (Maybe a)
    getHead :: m (Maybe a)
getHead = ([a] -> (Maybe a, [a])) -> m (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (([a] -> (Maybe a, [a])) -> m (Maybe a))
-> ([a] -> (Maybe a, [a])) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ \case
        (a
x:[a]
xs) -> (a -> Maybe a
forall a. a -> Maybe a
Just a
x, [a]
xs)
        [a]
xs -> (Maybe a
forall a. Maybe a
Nothing, [a]
xs)
    bit :: (MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) => m Bool
    bit :: m Bool
bit = m (Maybe Bool)
forall a (m :: * -> *). MonadState [a] m => m (Maybe a)
getHead m (Maybe Bool) -> (Maybe Bool -> m Bool) -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m Bool -> (Bool -> m Bool) -> Maybe Bool -> m Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m Bool
nextBit Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure where
        nextBit :: m Bool
nextBit = do
            Int
n <- Char -> Int
digitToInt (Char -> Int) -> m Char -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
hexDigitChar
            Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
3 Bool -> m () -> m Bool
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [Bool] -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
n (Int -> Bool) -> [Int] -> [Bool]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
2, Int
1, Int
0])
    bits :: (Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) => Int -> m a
    bits :: Int -> m a
bits Int
n = (a -> Bool -> a) -> a -> [Bool] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> Bool -> a
forall a. Num a => a -> Bool -> a
f a
0 ([Bool] -> a) -> m [Bool] -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m Bool -> m [Bool]
forall (m :: * -> *) a. Monad m => Int -> m a -> m [a]
count Int
n m Bool
forall (m :: * -> *) e s.
(MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m Bool
bit where f :: a -> Bool -> a
f a
acc Bool
b = a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ if Bool
b then a
1 else a
0
    literal :: (Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) => m a
    literal :: m a
literal = a -> m a
forall e s (m :: * -> *) b.
(MonadParsec e s m, Num b, MonadState [Bool] m, Token s ~ Char) =>
b -> m b
literal' a
0 where
        literal' :: b -> m b
literal' b
n = do
            Bool
b <- m Bool
forall (m :: * -> *) e s.
(MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m Bool
bit
            b
n' <- b -> b -> b
forall a. Num a => a -> a -> a
f b
n (b -> b) -> m b -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m b
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m a
bits Int
4
            if Bool
b then b -> m b
literal' b
n' else b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
n'
        f :: a -> a -> a
f a
acc a
n = a
16 a -> a -> a
forall a. Num a => a -> a -> a
* a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
n
    getOffsetBits :: (MonadState [a] m, MonadParsec e s m) => m Int
    getOffsetBits :: m Int
getOffsetBits = (-) (Int -> Int -> Int) -> (Int -> Int) -> Int -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4) (Int -> Int -> Int) -> m Int -> m (Int -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Int
forall e s (m :: * -> *). MonadParsec e s m => m Int
getOffset m (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([a] -> Int) -> m Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
    parseSpan :: (Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) => Int -> m [Packet a]
    parseSpan :: Int -> m [Packet a]
parseSpan Int
0 = [Packet a] -> m [Packet a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    parseSpan Int
n = do
        Int
offset0 <- m Int
forall a (m :: * -> *) e s.
(MonadState [a] m, MonadParsec e s m) =>
m Int
getOffsetBits
        Packet a
p <- m (Packet a)
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
m (Packet a)
packet
        Int
offset1 <- m Int
forall a (m :: * -> *) e s.
(MonadState [a] m, MonadParsec e s m) =>
m Int
getOffsetBits
        (:) Packet a
p ([Packet a] -> [Packet a]) -> m [Packet a] -> m [Packet a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m [Packet a]
forall a (m :: * -> *) e s.
(Num a, MonadState [Bool] m, MonadParsec e s m, Token s ~ Char) =>
Int -> m [Packet a]
parseSpan (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
offset1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offset0))

day16a :: Text -> Either (ParseErrorBundle Text Void) Int
day16a :: Text -> Either (ParseErrorBundle Text Void) Int
day16a Text
input = do
    Packet Int
packet <- Parsec Void Text (Packet Int)
-> String
-> Text
-> Either (ParseErrorBundle Text Void) (Packet Int)
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse (forall a e s (m :: * -> *).
(Num a, MonadParsec e s m, Token s ~ Char) =>
m (Packet a)
forall e s (m :: * -> *).
(Num Int, MonadParsec e s m, Token s ~ Char) =>
m (Packet Int)
parser @Int) String
"" Text
input
    Int -> Either (ParseErrorBundle Text Void) Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Either (ParseErrorBundle Text Void) Int)
-> Int -> Either (ParseErrorBundle Text Void) Int
forall a b. (a -> b) -> a -> b
$ Packet Int -> Int
forall a. Packet a -> Int
sumVersions Packet Int
packet
  where
    sumVersions :: Packet a -> Int
sumVersions Literal { Int
version :: Int
version :: forall a. Packet a -> Int
version } = Int
version
    sumVersions Operator { Int
version :: Int
version :: forall a. Packet a -> Int
version, [Packet a]
packets :: [Packet a]
packets :: forall a. Packet a -> [Packet a]
packets } = Int
version Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Packet a -> Int
sumVersions (Packet a -> Int) -> [Packet a] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Packet a]
packets)

day16b :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day16b :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day16b Text
input = do
    Packet Int
packet <- Parsec Void Text (Packet Int)
-> String
-> Text
-> Either (ParseErrorBundle Text Void) (Packet Int)
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse Parsec Void Text (Packet Int)
forall a e s (m :: * -> *).
(Num a, MonadParsec e s m, Token s ~ Char) =>
m (Packet a)
parser String
"" Text
input
    Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int))
-> Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Packet Int -> Maybe Int
forall a. (Num a, Ord a) => Packet a -> Maybe a
eval Packet Int
packet
  where
    eval :: Packet a -> Maybe a
eval Literal { a
value :: a
value :: forall a. Packet a -> a
value } = a -> Maybe a
forall a. a -> Maybe a
Just a
value
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
0, [Packet a]
packets :: [Packet a]
packets :: forall a. Packet a -> [Packet a]
packets } = [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([a] -> a) -> Maybe [a] -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Packet a -> Maybe a) -> [Packet a] -> Maybe [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Packet a -> Maybe a
eval [Packet a]
packets
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
1, [Packet a]
packets :: [Packet a]
packets :: forall a. Packet a -> [Packet a]
packets } = [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([a] -> a) -> Maybe [a] -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Packet a -> Maybe a) -> [Packet a] -> Maybe [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Packet a -> Maybe a
eval [Packet a]
packets
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
2, [Packet a]
packets :: [Packet a]
packets :: forall a. Packet a -> [Packet a]
packets } = [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([a] -> a) -> Maybe [a] -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Packet a -> Maybe a) -> [Packet a] -> Maybe [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Packet a -> Maybe a
eval [Packet a]
packets
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
3, [Packet a]
packets :: [Packet a]
packets :: forall a. Packet a -> [Packet a]
packets } = [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([a] -> a) -> Maybe [a] -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Packet a -> Maybe a) -> [Packet a] -> Maybe [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Packet a -> Maybe a
eval [Packet a]
packets
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
5, packets :: forall a. Packet a -> [Packet a]
packets = [Packet a -> Maybe a
eval -> Just a
lhs, Packet a -> Maybe a
eval -> Just a
rhs] } =
        a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ if a
lhs a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
rhs then a
1 else a
0
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
6, packets :: forall a. Packet a -> [Packet a]
packets = [Packet a -> Maybe a
eval -> Just a
lhs, Packet a -> Maybe a
eval -> Just a
rhs] } =
        a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ if a
lhs a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
rhs then a
1 else a
0
    eval Operator { tag :: forall a. Packet a -> Int
tag = Int
7, packets :: forall a. Packet a -> [Packet a]
packets = [Packet a -> Maybe a
eval -> Just a
lhs, Packet a -> Maybe a
eval -> Just a
rhs] } =
        a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ if a
lhs a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
rhs then a
1 else a
0
    eval Packet a
_ = Maybe a
forall a. Maybe a
Nothing