Tree Insert using State Monad

Question

I have a tree and insert operation defined as in "Learn You a Haskell for Great Good!" :

data Tree a = EmptyTree | Node a (Tree a) (Tree a) deriving (Show, Read, Eq) 

treeInsert :: (Ord a) => a -> Tree a -> Tree a  
treeInsert x EmptyTree = Node x EmptyTree EmptyTree
treeInsert x (Node a left right)   
    | x == a = Node x left right  
    | x < a  = Node a (treeInsert x left) right  
    | x > a  = Node a left (treeInsert x right)   

I would like to reimplement treeInsert using State Monad, but I'm not even sure how function declaration should look like. I have this so far:

treeInsert :: (Ord a) => a -> Tree a -> State (Tree a) a

How would you write treeInsert using State Monad?

Answer 1

Warning: This answer contains spoilers.

You can fairly easily write a wrapper around your existing treeInsert function that allows you do use do-notation the way you want. As per the comments, there's a function modify that takes a modifying function f :: s -> s and turns it into a State s () which is an "action" to modify a state s. That means you can write:

stateTreeInsert :: (Ord a) => a -> State (Tree a) ()
stateTreeInsert x = modify (treeInsert x)

or more succinctly:

stateTreeInsert :: (Ord a) => a -> State (Tree a) ()
stateTreeInsert = modify . treeInsert

Then, you can define an "action" like:

insertSomeStuff :: (Ord a, Num a) => State (Tree a) ()
insertSomeStuff = do
  stateTreeInsert 0
  stateTreeInsert 1
  stateTreeInsert 2

and then apply it to a particular tree using execState:

main = print $ execState insertSomeStuff EmptyTree

However, I guess you were more interested in re-implementing treeInsert from scratch in a state-manipulating form.

The problem is that the "straightforward" way of doing this isn't very interesting or idiomatic. It's just awkward. It would look something like this:

awkwardTreeInsert :: (Ord a) => a -> State (Tree a) ()
awkwardTreeInsert x = do
  t <- get
  case t of
    EmptyTree -> put $ Node x EmptyTree EmptyTree
    Node a l r -> case compare x a of
      LT -> do put l                 -- replace tree with left child
               awkwardTreeInsert x   -- insert into left child
               l' <- get             -- get the answer
               put $ Node a l' r     -- overwrite with whole tree w/ updated left child
      GT -> do put r
               awkwardTreeInsert x
               r' <- get
               put $ Node a l r'
      EQ -> return ()

The issue here is that the state, as we've written it, can only hold one tree at once. So, if we want to call the algorithm recursively to insert something into a branch, we need to overwrite the "big tree" with one of its children, run the recursive insertion, get the answer, and overwrite it with the "big tree" with the appropriate child replaced.

Anyway, it works the same way as stateTreeInsert so:

insertSomeStuff :: (Ord a, Num a) => State (Tree a) ()
insertSomeStuff = do
  awkwardTreeInsert 0
  awkwardTreeInsert 1
  awkwardTreeInsert 2

main = print $ execState insertSomeStuff EmptyTree