{-# Language PatternSynonyms, DeriveTraversable, QuasiQuotes, BlockArguments, LambdaCase, ImportQualifiedPost #-}
{-|
Module      : Main
Description : Day 21 solution
Copyright   : (c) Eric Mertens, 2022
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2022/day/21>

>>> :{
:main +
  "root: pppw + sjmn\n\
  \dbpl: 5\n\
  \cczh: sllz + lgvd\n\
  \zczc: 2\n\
  \ptdq: humn - dvpt\n\
  \dvpt: 3\n\
  \lfqf: 4\n\
  \humn: 5\n\
  \ljgn: 2\n\
  \sjmn: drzm * dbpl\n\
  \sllz: 4\n\
  \pppw: cczh / lfqf\n\
  \lgvd: ljgn * ptdq\n\
  \drzm: hmdt - zczc\n\
  \hmdt: 32\n"
:}
152
301

-}
module Main where

import Data.Functor ((<&>))
import Data.Map (Map)
import Data.Map qualified as Map

import Advent (format)
import Advent.Fix (Fix(Fix))

-- |
-- >>> :main
-- 110181395003396
-- 3721298272959
main :: IO ()
IO ()
main =
 do Map String (Expr String)
input <- [(String, Either Int (String, Char, String))]
-> Map String (Expr String)
buildMap ([(String, Either Int (String, Char, String))]
 -> Map String (Expr String))
-> IO [(String, Either Int (String, Char, String))]
-> IO (Map String (Expr String))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2022 21 (%s: (%d|%s %c %s)%n)*|]

    -- part 1
    case Map String (Expr String) -> Fix Expr
evalRoot Map String (Expr String)
input of
      I Int
p1 -> Int -> IO ()
forall a. Show a => a -> IO ()
print Int
p1
      Fix Expr
_ -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unsupported part 1"

    -- part 2
    case Map String (Expr String) -> Fix Expr
evalRoot (String
-> Expr String
-> Map String (Expr String)
-> Map String (Expr String)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
"humn" Expr String
forall a. Expr a
Answer Map String (Expr String)
input) of
      Fix (Add Fix Expr
x (I Int
y)) -> Int -> IO ()
forall a. Show a => a -> IO ()
print (Fix Expr -> Int -> Int
equal Fix Expr
x Int
y)
      Fix Expr
_ -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unsupported part 2"

evalRoot :: Map String (Expr String) -> Fix Expr
evalRoot :: Map String (Expr String) -> Fix Expr
evalRoot Map String (Expr String)
env = Map String (Fix Expr)
env' Map String (Fix Expr) -> String -> Fix Expr
forall k a. Ord k => Map k a -> k -> a
Map.! String
"root"
  where env' :: Map String (Fix Expr)
env' = (Expr String -> Fix Expr)
-> Map String (Expr String) -> Map String (Fix Expr)
forall a b. (a -> b) -> Map String a -> Map String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Fix Expr -> Fix Expr
constProp (Fix Expr -> Fix Expr)
-> (Expr String -> Fix Expr) -> Expr String -> Fix Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map String (Fix Expr) -> Expr String -> Fix Expr
forall (f :: * -> *) a.
(Functor f, Ord a) =>
Map a (Fix f) -> f a -> Fix f
tie Map String (Fix Expr)
env') Map String (Expr String)
env

tie :: (Functor f, Ord a) => Map a (Fix f) -> f a -> Fix f
tie :: forall (f :: * -> *) a.
(Functor f, Ord a) =>
Map a (Fix f) -> f a -> Fix f
tie Map a (Fix f)
m f a
e = f (Fix f) -> Fix f
forall (f :: * -> *). f (Fix f) -> Fix f
Fix ((a -> Fix f) -> f a -> f (Fix f)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map a (Fix f)
m Map a (Fix f) -> a -> Fix f
forall k a. Ord k => Map k a -> k -> a
Map.!) f a
e)

buildMap :: [(String, Either Int (String, Char, String))] -> Map String (Expr String)
buildMap :: [(String, Either Int (String, Char, String))]
-> Map String (Expr String)
buildMap [(String, Either Int (String, Char, String))]
xs =
  [(String, Either Int (String, Char, String))]
-> Map String (Either Int (String, Char, String))
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(String, Either Int (String, Char, String))]
xs Map String (Either Int (String, Char, String))
-> (Either Int (String, Char, String) -> Expr String)
-> Map String (Expr String)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    Left Int
i          -> Int -> Expr String
forall a. Int -> Expr a
Lit Int
i
    Right (String
a,Char
'*',String
b) -> String -> String -> Expr String
forall a. a -> a -> Expr a
Mul String
a String
b
    Right (String
a,Char
'+',String
b) -> String -> String -> Expr String
forall a. a -> a -> Expr a
Add String
a String
b
    Right (String
a,Char
'/',String
b) -> String -> String -> Expr String
forall a. a -> a -> Expr a
Div String
a String
b
    Right (String
a,Char
'-',String
b) -> String -> String -> Expr String
forall a. a -> a -> Expr a
Sub String
a String
b
    Either Int (String, Char, String)
_               -> String -> Expr String
forall a. HasCallStack => String -> a
error String
"bad expression"

constProp :: Fix Expr -> Fix Expr
constProp :: Fix Expr -> Fix Expr
constProp = \case
  Fix (Add (I Int
x) (I Int
y)) -> Int -> Fix Expr
I (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y)
  Fix (Sub (I Int
x) (I Int
y)) -> Int -> Fix Expr
I (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
y)
  Fix (Mul (I Int
x) (I Int
y)) -> Int -> Fix Expr
I (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
y)
  Fix (Div (I Int
x) (I Int
y)) | (Int
z,Int
0) <- Int
x Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
y -> Int -> Fix Expr
I Int
z
  Fix Expr
e -> Fix Expr
e

equal :: Fix Expr -> Int -> Int
equal :: Fix Expr -> Int -> Int
equal (Fix (Div Fix Expr
x (I Int
y))) Int
z = Fix Expr -> Int -> Int
equal Fix Expr
x (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
z)
equal (Fix (Add (I Int
x) Fix Expr
y)) Int
z = Fix Expr -> Int -> Int
equal Fix Expr
y (Int
zInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
x)
equal (Fix (Add Fix Expr
x (I Int
y))) Int
z = Fix Expr -> Int -> Int
equal Fix Expr
x (Int
zInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
y)
equal (Fix (Mul (I Int
x) Fix Expr
y)) Int
z | (Int
z',Int
0) <- Int
z Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
x = Fix Expr -> Int -> Int
equal Fix Expr
y Int
z'
equal (Fix (Mul Fix Expr
x (I Int
y))) Int
z | (Int
z',Int
0) <- Int
z Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
y = Fix Expr -> Int -> Int
equal Fix Expr
x Int
z'
equal (Fix (Sub (I Int
x) Fix Expr
y)) Int
z = Fix Expr -> Int -> Int
equal Fix Expr
y (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
z)
equal (Fix (Sub Fix Expr
x (I Int
y))) Int
z = Fix Expr -> Int -> Int
equal Fix Expr
x (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
z)
equal (Fix Expr (Fix Expr)
Answer) Int
x = Int
x
equal Fix Expr
_ Int
_ = String -> Int
forall a. HasCallStack => String -> a
error String
"stuck"

pattern I :: Int -> Fix Expr
pattern $mI :: forall {r}. Fix Expr -> (Int -> r) -> ((# #) -> r) -> r
$bI :: Int -> Fix Expr
I i = Fix (Lit i)

data Expr a
    = Add a a
    | Sub a a
    | Mul a a
    | Div a a
    | Lit Int
    | Answer
    deriving (Int -> Expr a -> String -> String
[Expr a] -> String -> String
Expr a -> String
(Int -> Expr a -> String -> String)
-> (Expr a -> String)
-> ([Expr a] -> String -> String)
-> Show (Expr a)
forall a. Show a => Int -> Expr a -> String -> String
forall a. Show a => [Expr a] -> String -> String
forall a. Show a => Expr a -> String
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Expr a -> String -> String
showsPrec :: Int -> Expr a -> String -> String
$cshow :: forall a. Show a => Expr a -> String
show :: Expr a -> String
$cshowList :: forall a. Show a => [Expr a] -> String -> String
showList :: [Expr a] -> String -> String
Show, (forall a b. (a -> b) -> Expr a -> Expr b)
-> (forall a b. a -> Expr b -> Expr a) -> Functor Expr
forall a b. a -> Expr b -> Expr a
forall a b. (a -> b) -> Expr a -> Expr b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Expr a -> Expr b
fmap :: forall a b. (a -> b) -> Expr a -> Expr b
$c<$ :: forall a b. a -> Expr b -> Expr a
<$ :: forall a b. a -> Expr b -> Expr a
Functor)