import Control.Monad.Trans.State.Lazy import Control.Monad (foldM) data Input = Plus | Minus | ToggleNegate | ToggleEnabled type Emission = Integer type Accum = [Emission] type Output = [Emission] type Negated = Bool type Enabled = Bool toInput :: Char -> Input toInput '+' = Plus toInput '-' = Minus toInput 'n' = ToggleNegate toInput 'e' = ToggleEnabled toInput _ = error "Invalid input representation" -- Determine new state of state machine along with transition emission step :: (Negated, Enabled, Input) -> (Negated, Enabled, Emission) step (n, e, ToggleNegate) = (not n, e, 0) step (n, e, ToggleEnabled) = (n, not e, 0) step (n, False, i) = (n, False, 0) step (n, e, Plus) = (n, e, if n then -1 else 1) step (n, e, Minus) = (n, e, if n then 1 else -1) -- Helper function for "evaluate"'s foldM mapEmissions :: Accum -> Input -> State (Negated, Enabled) Output mapEmissions accum input = do (negated, enabled) <- get let (negated', enabled', emission) = step (negated, enabled, input) put (negated', enabled') return (accum ++ [emission]) -- Process an input string and return the result -- (False, True) is the start state: (not negated, enabled) evaluate :: [Input] -> Output evaluate inputs = evalState (foldM mapEmissions [] inputs) (False, True) -- Convenience function for output formatting shouldEqual :: String -> Integer -> IO () shouldEqual input expected = do let actual = (sum . evaluate . map toInput) input putStrLn $ "Expected " ++ show expected ++ ", got " ++ show actual ++ ": " ++ input main :: IO () main = do "+-n--n" `shouldEqual` 2 "+e----e++" `shouldEqual` 3 "-n++e++e--+-n++" `shouldEqual` 1