Kata: Mathematical AST

Few days ago, with the Software Crafters Lyon, we have tackled the Mathematical AST.

This code kata suggests tackling it as outside-in TDD.

Which is a massive bias for me, as I have learned, a long time ago, in the K&R to handle it with a stack.

The basic idea of the kata is to accept reverse polish notation, which is a parenthesis-less way to express mathematical expressions, as follows:

We will assume expressions are valid.

As such, we can start handling a simple number as an expression, as follows:

spec :: Spec
spec =
  describe "RPN Evaluator" $ do
    forM_
      [ ("0", 0),
        ("42", 42)
      ]
      $ \(expr, expected) ->
        it (show expr <> " should be evaluated to " <> show expected) $
          eval expr `shouldBe` expected

eval :: String -> Int
eval = read

The next could be as simple as handling sum, as follows:

-- ("1 2 +", 3)
-- ("2 2 +", 4)
eval :: String -> Int
eval expr =
  case words expr of
    [n] -> read n
    [x, y, _op] -> read x + read y

The next simplest step is to add the minus operation, as follows:

-- ("2 1 -", 1)
-- ("2 2 -", 0)
eval :: String -> Int
eval expr =
  case words expr of
    [n] -> read n
    [x, y, op] ->
      let f =
            case op of
              "+" -> (+)
              "-" -> (-)
       in read x `f` read y

The “ultimate” step is to handle multiple operations.

To proceed, we have to introduce a stack, last in first out, in which we will add numbers and, on operators, pop two values, apply them to the operator, and then push the result on the stack, as follows:

-- ("1 2 3 * +", 7),
-- ("1 2 + 3 *", 9),
-- ("1 2 3 * -", -5)
eval :: String -> Int
eval = head . foldl go [] . words
  where
    go stack =
      \case
        "+" -> x + y : zs
        "-" -> x - y : zs
        "*" -> x * y : zs
        n -> read n : stack
      where
        (y : x : zs) = stack

The evaluation of 1 2 3 * - can be visualized as follows:

Then, getting the evaluation result only takes to get the top of the stack.

It should be enough to fulfill the first part of this kata, but I would like to spend some time on the design of the stack.

Let's extract proper types and functions, as follows:

newtype Memory a = Memory [a]

emptyMemory :: Memory a
emptyMemory = Memory []

number :: a -> Memory a -> Memory a
number x (Memory xs) = Memory (x : xs)

operator :: (a -> a -> a) -> Memory a -> Memory a
operator f (Memory (y : x : zs)) = Memory (f x y : zs)

top :: Memory a -> a
top (Memory (x : _)) = x

Which simplifies our eval implementation as follows:

eval :: String -> Int
eval = top . foldl (flip go) emptyMemory . words
  where
    go :: String -> Memory Int -> Memory Int
    go =
      \case
        "+" -> operator (+)
        "-" -> operator (-)
        "*" -> operator (*)
        n -> number (read n)

During the session, I have asked myself, "Can this be done without a stack?”

Actually, we can borrow Church-encoded non-empty list/stream, as follows:

newtype Memory a
  = Memory
      ( forall b.
        (a -> Memory a -> b) ->
        b
      )

runMemory :: (a -> Memory a -> b) -> Memory a -> b
runMemory f (Memory m) = m f

Then, we should rewrite our helpers as follows:

emptyMemory :: Memory a
emptyMemory = Memory $ \f -> f empty emptyMemory
  where
    empty = error "empty memory"

number :: a -> Memory a -> Memory a
number x m = Memory $ \f -> f x m

operator :: (a -> a -> a) -> Memory a -> Memory a
operator f = runMemory $ \y -> runMemory (\x -> number (f x y))

top :: Memory a -> a
top = runMemory const

All of this, without touching eval.