Day21.hs

module Javran.AdventOfCode.Y2022.Day21 (

) where

import Control.Monad.RWS.CPS
import Data.Char
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Javran.AdventOfCode.Prelude
import Javran.AdventOfCode.Y2022.Day21.Types
import qualified Javran.AdventOfCode.Y2022.Day21.Z3 as Z3
import Text.ParserCombinators.ReadP hiding (count, get, many)

data Day21 deriving (Generic)

{-
  IMO it's a really questionable decision that a programming challenge would
  just using division without any specification.

  We'd better make sure.
 -}
checkedQuot :: V -> V -> V
checkedQuot a b
  | b <= 0 = error $ "divisor <= 0: " <> show b
  | a < 0 = error $ "dividend < 0" <> show a
  | r /= 0 = error $ "not evenly divisible: " <> show (a, b)
  | otherwise = q
  where
    (q, r) = quotRem a b

stmtP :: ReadP Stmt
stmtP = do
  let varP = T.pack <$> munch1 isAsciiLower
  lhs <- varP <* strP ": "
  rhs <-
    (ELit <$> decimal1P)
      <++ ( do
              vL <- varP <* charP ' '
              op <-
                (Plus <$ strP "+ ")
                  <++ (Minus <$ strP "- ")
                  <++ (Mult <$ strP "* ")
                  <++ (Div <$ strP "/ ")
              vR <- varP
              pure $ EOp op vL vR
          )
  pure (lhs, rhs)

-- Evaluation context for both parts.
type M v =
  RWS
    (M.Map T.Text Expr) -- Known statements
    () -- Unused Writer
    (M.Map T.Text v) -- Cache of evaluated variables

opToFn :: Op -> V -> V -> V
opToFn = \case
  Plus -> (+)
  Minus -> (-)
  Mult -> (*)
  Div -> checkedQuot

eval :: T.Text -> M V V
eval vn =
  gets (M.!? vn) >>= \case
    Just v -> pure v
    Nothing -> do
      e <- asks (M.! vn)
      val <- case e of
        ELit v -> pure v
        EOp op l r ->
          opToFn op <$> eval l <*> eval r
      val <$ modify (M.insert vn val)

{-
  Represents a partial value that has exactly one hole
  (an unkonwn variable) in it.

  The idea is that whenever we encounter a hole,
  we only evaluate the part that we can evaluate and
  builds a larger hole in return.

  By doing so, we gradually build up a stack of constraints
  that we can tear down layer-by-layer once we know what this
  value is supposed to be evaluated to.
 -}
data Partial
  = PHole
  | PHoleInL Partial Op V
  | PHoleInR V Op Partial
  deriving (Show)

type V' = Either Partial V

{-
  Evaluates whatever we can evaluate, and builds up the partial value.
  (see comment of `Partial` for more details)
 -}
evalP2 :: T.Text -> M V' V'
evalP2 "humn" = pure $ Left PHole
evalP2 vn =
  gets (M.!? vn) >>= \case
    Just v -> pure v
    Nothing -> do
      e <- asks (M.! vn)
      val <- case e of
        ELit v -> pure (Right v)
        EOp op l r -> do
          vL <- evalP2 l
          vR <- evalP2 r
          pure case (vL, vR) of
            (Left _, Left _) ->
              error "cannot handle both sides being partial"
            (Left p, Right r') ->
              Left $ PHoleInL p op r'
            (Right l', Left p) ->
              Left $ PHoleInR l' op p
            (Right l', Right r') ->
              Right $ opToFn op l' r'
      val <$ modify (M.insert vn val)

{-
  Constraint solving at its simplest form.

  The `{}` below notes about where the hole is supposed to be.
 -}
solvePartial :: V -> Partial -> V
solvePartial = fix \go v -> \case
  PHole -> v
  PHoleInL p' op r ->
    case op of
      Plus ->
        -- {} + r = v
        go (v - r) p'
      Minus ->
        -- {} - r = v
        go (v + r) p'
      Mult ->
        -- {} * r = v
        go (dv v r) p'
      Div ->
        -- {} / r = v
        go (v * r) p'
  PHoleInR l op p' ->
    case op of
      Plus ->
        -- l + {} = v
        go (v - l) p'
      Minus ->
        -- l - {} = v
        go (l - v) p'
      Mult ->
        -- l * {} = v
        go (dv v l) p'
      Div ->
        -- l / {} = v
        go (dv l v) p'
  where
    dv a b =
      let (q, r) = quotRem a b
       in if r == 0
            then q
            else error $ "Not evenly divisible: " <> show (a, b)

instance Solution Day21 where
  solutionRun _ SolutionContext {getInputS, answerShow} = do
    stmts <-
      M.fromList
        . consumeOrDie (sepBy1 stmtP (charP '\n') <* skipSpaces)
        <$> getInputS
    do
      let (ans, _, _) = runRWS (eval "root") stmts M.empty
      answerShow ans
    let solveWithZ3 = False
    if solveWithZ3
      then do
        ans <- Z3.solve stmts
        answerShow ans
      else do
        let stmts' = M.delete "humn" . M.delete "root" $ stmts
            (EOp _ lhs rhs) = stmts M.! "root"
            (lrVals, _, _) =
              runRWS
                ((,) <$> evalP2 lhs <*> evalP2 rhs)
                stmts'
                M.empty
        answerShow case lrVals of
          (Left p, Right v) -> solvePartial v p
          (Right v, Left p) -> solvePartial v p
          (Left _, Left _) -> error "both sides are partial"
          (Right _, Right _) -> error "nothing to solve"