-- start snippet imports
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StrictData #-}
module Main where
import qualified Data.Array as Array
import Data.Char (digitToInt, isAlpha, isDigit)
import Data.Foldable (foldl', foldr')
import Data.Function (fix)
import Data.Functor (($>))
import qualified Data.Graph as Graph
import Data.List (intercalate, (\\))
import qualified Data.Map.Strict as Map
import System.Environment (getArgs)
import qualified Text.ParserCombinators.ReadP as P
import Prelude hiding (GT, LT)
-- end snippet imports
-- start snippet parts-parser
data Part = Part
{ partX :: Int,
partM :: Int,
partA :: Int,
partS :: Int
} deriving (Show)
data Rating = X | M | A | S deriving (Show, Eq)
emptyPart :: Part
emptyPart = Part 0 0 0 0
addRating :: Part -> (Rating, Int) -> Part
addRating p (r, v) = case r of
X -> p {partX = v}
M -> p {partM = v}
A -> p {partA = v}
S -> p {partS = v}
partParser :: P.ReadP Part
partParser =
foldl' addRating emptyPart
<$> P.between (P.char '{') (P.char '}')
(partRatingParser `P.sepBy1` P.char ',')
partRatingParser :: P.ReadP (Rating, Int)
partRatingParser =
(,) <$> ratingParser <*> (P.char '=' *> intParser)
ratingParser :: P.ReadP Rating
ratingParser =
P.get >>= \case
'x' -> pure X
'm' -> pure M
'a' -> pure A
's' -> pure S
_ -> P.pfail
intParser :: P.ReadP Int
intParser =
foldl' (\n d -> n * 10 + d) 0 <$> P.many1 digitParser
digitParser :: P.ReadP Int
digitParser = digitToInt <$> P.satisfy isDigit
parse :: (Show a) => P.ReadP a -> String -> Either String a
parse parser text = case P.readP_to_S (parser <* P.eof) text of
[(res, "")] -> Right res
[(_, s)] -> Left $ "Leftover input: " <> s
out -> Left $ "Unexpected output: " <> show out
-- end snippet parts-parser
-- start snippet system
newtype System =
System (Map.Map WorkflowName Workflow)
deriving (Show, Eq)
data Workflow = Workflow
{ wName :: WorkflowName,
wRules :: [Rule]
} deriving (Show, Eq)
type WorkflowName = String
data Rule
= AtomicRule AtomicRule
| If Condition AtomicRule
deriving (Show, Eq)
data AtomicRule
= Jump WorkflowName
| Accept
| Reject
deriving (Show, Eq, Ord)
data Condition
= Comparison Rating CmpOp Int
deriving (Show, Eq)
data CmpOp = LT | GT deriving (Show, Eq)
-- end snippet system
-- start snippet system-parser
systemParser :: P.ReadP System
systemParser =
System
. foldl' (\m wf -> Map.insert (wName wf) wf m) Map.empty
<$> workflowParser `P.endBy1` P.char '\n'
workflowParser :: P.ReadP Workflow
workflowParser =
Workflow
<$> P.many1 (P.satisfy isAlpha)
<*> P.between (P.char '{') (P.char '}')
(ruleParser `P.sepBy1` P.char ',')
ruleParser :: P.ReadP Rule
ruleParser =
(AtomicRule <$> atomicRuleParser) P.<++ ifRuleParser
ifRuleParser :: P.ReadP Rule
ifRuleParser =
If
<$> (Comparison <$> ratingParser <*> cmpOpParser <*> intParser)
<*> (P.char ':' *> atomicRuleParser)
atomicRuleParser :: P.ReadP AtomicRule
atomicRuleParser = do
c : _ <- P.look
case c of
'A' -> P.char 'A' $> Accept
'R' -> P.char 'R' $> Reject
_ -> (Jump .) . (:) <$> P.char c <*> P.many1 (P.satisfy isAlpha)
cmpOpParser :: P.ReadP CmpOp
cmpOpParser = P.choice [P.char '<' $> LT, P.char '>' $> GT]
-- end snippet system-parser
-- start snippet input-parser
data Input = Input System [Part] deriving (Show)
inputParser :: P.ReadP Input
inputParser =
Input
<$> systemParser
<*> (P.char '\n' *> partParser `P.endBy1` P.char '\n')
-- end snippet input-parser
-- start snippet interpreter
runSystem :: System -> Part -> Bool
runSystem (System system) part = runRule $ Jump "in"
where
runRule = \case
Accept -> True
Reject -> False
Jump wfName -> jump wfName
jump wfName = case Map.lookup wfName system of
Just workflow -> runRules $ wRules workflow
Nothing ->
error $ "Workflow not found in system: " <> wfName
runRules = \case
(rule : rest) -> case rule of
AtomicRule aRule -> runRule aRule
If cond aRule ->
if evalCond cond
then runRule aRule
else runRules rest
_ -> error "Workflow ended without accept/reject"
evalCond = \case
Comparison r LT value -> rating r < value
Comparison r GT value-> rating r > value
rating = \case
X -> partX part
M -> partM part
A -> partA part
S -> partS part
-- end snippet interpreter
-- start snippet solve
solve :: Input -> Int
solve (Input system parts) =
sum
. map (\(Part x m a s) -> x + m + a + s)
. filter (runSystem system)
$ parts
-- end snippet solve
-- start snippet graph-algos
type Graph' a =
(Graph.Graph, Graph.Vertex -> (a, [a]), a -> Maybe Graph.Vertex)
cfGraph :: Map.Map WorkflowName Workflow -> Graph' WorkflowName
cfGraph system =
graphFromMap
. Map.toList
. flip Map.map system
$ \(Workflow _ rules) ->
flip concatMap rules $ \case
AtomicRule (Jump wfName) -> [wfName]
If _ (Jump wfName) -> [wfName]
_ -> []
where
graphFromMap :: (Ord a) => [(a, [a])] -> Graph' a
graphFromMap m =
let (graph, nLookup, vLookup) =
Graph.graphFromEdges $ map (\(f, ts) -> (f, f, ts)) m
in (graph, \v -> let (x, _, xs) = nLookup v in (x, xs), vLookup)
toposortWorkflows :: Map.Map WorkflowName Workflow -> [WorkflowName]
toposortWorkflows system =
let (cfg, nLookup, _) = cfGraph system
in map (fst . nLookup) $ Graph.topSort cfg
-- end snippet graph-algos
-- start snippet compiler
class ToC a where
toC :: a -> String
instance ToC Part where
toC (Part x m a s) =
"{" <> intercalate ", " (map show [x, m, a, s]) <> "}"
instance ToC CmpOp where
toC = \case
LT -> "<"
GT -> ">"
instance ToC Rating where
toC = \case
X -> "x"
M -> "m"
A -> "a"
S -> "s"
instance ToC AtomicRule where
toC = \case
Accept -> "return true;"
Reject -> "return false;"
Jump wfName -> "goto " <> wfName <> ";"
instance ToC Condition where
toC = \case
Comparison rating op val ->
toC rating <> " " <> toC op <> " " <> show val
instance ToC Rule where
toC = \case
AtomicRule aRule -> toC aRule
If cond aRule ->
"if (" <> toC cond <> ") { " <> toC aRule <> " }"
instance ToC Workflow where
toC (Workflow wfName rules) =
wfName
<> ":\n"
<> intercalate "\n" (map ((" " <>) . toC) rules)
instance ToC System where
toC (System system) =
intercalate
"\n"
[ "bool runSystem(int x, int m, int a, int s) {",
" goto in;",
intercalate
"\n"
(map (toC . (system Map.!)) $ toposortWorkflows system),
"}"
]
instance ToC Input where
toC (Input system parts) =
intercalate
"\n"
[ "#include <stdbool.h>",
"#include <stdio.h>\n",
toC system,
"int main() {",
" int parts[][4] = {",
intercalate ",\n" (map ((" " <>) . toC) parts),
" };",
" int totalRating = 0;",
" for(int i = 0; i < " <> show (length parts) <> "; i++) {",
" int x = parts[i][0];",
" int m = parts[i][1];",
" int a = parts[i][2];",
" int s = parts[i][3];",
" if (runSystem(x, m, a, s)) {",
" totalRating += x + m + a + s;",
" }",
" }",
" printf(\"%d\", totalRating);",
" return 0;",
"}"
]
-- end snippet compiler
-- start snippet simplify-workflows
simplifyWorkflows :: System -> System
simplifyWorkflows (System system) =
System $ Map.map simplifyWorkflow system
where
simplifyWorkflow (Workflow name rules) =
Workflow name
$ foldr'
( \r rs -> case rs of
[r'] | ruleOutcome r == ruleOutcome r' -> rs
_ -> r : rs
)
[last rules]
$ init rules
ruleOutcome = \case
If _ aRule -> aRule
AtomicRule aRule -> aRule
-- end snippet simplify-workflows
-- start snippet inline-redundant-jumps
inlineRedundantJumps :: System -> System
inlineRedundantJumps (System system) =
System $
foldl' (flip Map.delete) (Map.map inlineJumps system) $
Map.keys redundantJumps
where
redundantJumps =
Map.map (\wf -> let ~(AtomicRule rule) = head $ wRules wf in rule)
. Map.filter (\wf -> length (wRules wf) == 1)
$ system
inlineJumps (Workflow name rules) =
Workflow name $ map inlineJump rules
inlineJump = \case
AtomicRule (Jump wfName)
| Map.member wfName redundantJumps ->
AtomicRule $ redundantJumps Map.! wfName
If cond (Jump wfName)
| Map.member wfName redundantJumps ->
If cond $ redundantJumps Map.! wfName
rule -> rule
-- end snippet inline-redundant-jumps
-- start snippet remove-jumps
removeJumps :: System -> System
removeJumps (System system) =
let system' =
foldl' (flip $ Map.adjust removeJumpsWithSingleJumper) system $
toposortWorkflows system
in System
. foldl' (flip Map.delete) system'
. (\\ ["in"])
$ workflowsWithNJumpers 0 system'
where
removeJumpsWithSingleJumper (Workflow name rules) =
Workflow name $
init rules <> case last rules of
AtomicRule (Jump wfName)
| wfName `elem` workflowsWithSingleJumper ->
let (Workflow _ rules') = system Map.! wfName
in rules'
rule -> [rule]
workflowsWithSingleJumper = workflowsWithNJumpers 1 system
workflowsWithNJumpers n sys =
let (cfg, nLookup, _) = cfGraph sys
in map (fst . nLookup . fst)
. filter (\(_, d) -> d == n)
. Array.assocs
. Graph.indegree
$ cfg
-- end snippet remove-jumps
-- start snippet optimize
optimize :: System -> System
optimize =
applyTillUnchanged
(removeJumps . inlineRedundantJumps . simplifyWorkflows)
where
applyTillUnchanged :: (Eq a) => (a -> a) -> a -> a
applyTillUnchanged f =
fix (\recurse x -> if f x == x then x else recurse (f x))
-- end snippet optimize
-- start snippet main
main :: IO ()
main = do
file <- head <$> getArgs
code <- readFile file
case parse inputParser code of
Right (Input system parts) ->
putStrLn . toC $ Input (optimize system) parts
Left err -> error err
-- end snippet main