fork download
  1. import Control.Monad.Trans.State.Lazy
  2. import Control.Monad (foldM)
  3.  
  4. data Input = Plus | Minus | ToggleNegate | ToggleEnabled
  5. type Emission = Integer
  6. type Accum = [Emission]
  7. type Output = [Emission]
  8. type Negated = Bool
  9. type Enabled = Bool
  10.  
  11. toInput :: Char -> Input
  12. toInput '+' = Plus
  13. toInput '-' = Minus
  14. toInput 'n' = ToggleNegate
  15. toInput 'e' = ToggleEnabled
  16. toInput _ = error "Invalid input representation"
  17.  
  18. -- Determine new state of state machine along with transition emission
  19. step :: (Negated, Enabled, Input) -> (Negated, Enabled, Emission)
  20. step (n, e, ToggleNegate) = (not n, e, 0)
  21. step (n, e, ToggleEnabled) = (n, not e, 0)
  22. step (n, False, i) = (n, False, 0)
  23. step (n, e, Plus) = (n, e, if n then -1 else 1)
  24. step (n, e, Minus) = (n, e, if n then 1 else -1)
  25.  
  26. -- Helper function for "evaluate"'s foldM
  27. mapEmissions :: Accum -> Input -> State (Negated, Enabled) Output
  28. mapEmissions accum input = do
  29. (negated, enabled) <- get
  30. let (negated', enabled', emission) = step (negated, enabled, input)
  31. put (negated', enabled')
  32. return (accum ++ [emission])
  33.  
  34. -- Process an input string and return the result
  35. -- (False, True) is the start state: (not negated, enabled)
  36. evaluate :: [Input] -> Output
  37. evaluate inputs = evalState (foldM mapEmissions [] inputs) (False, True)
  38.  
  39.  
  40. -- Convenience function for output formatting
  41. shouldEqual :: String -> Integer -> IO ()
  42. shouldEqual input expected = do
  43. let actual = (sum . evaluate . map toInput) input
  44. putStrLn $ "Expected " ++ show expected ++ ", got " ++ show actual ++ ": " ++ input
  45.  
  46. main :: IO ()
  47. main = do
  48. "+-n--n" `shouldEqual` 2
  49. "+e----e++" `shouldEqual` 3
  50. "-n++e++e--+-n++" `shouldEqual` 1
Success #stdin #stdout 0s 4164KB
stdin
Standard input is empty
stdout
Expected 2, got 2: +-n--n
Expected 3, got 3: +e----e++
Expected 1, got 1: -n++e++e--+-n++