-- nondeterministic state monad
module State where

data State s a = ST (s -> [(a,s)])

instance Monad (State s) where
  return x = ST (\s -> [(x,s)])
  fail   _ = ST (\_ -> [])

  ST a >>= f = ST (\s -> [ res | (x,s') <- a s, let ST b = f x, res <- b s' ])

runState :: s -> State s a -> [a]
runState s (ST a) = map fst (a s)

update :: (s -> s) -> State s s
update upd = ST (\s -> [(s,upd s)])

get :: State s s
get = update id

set :: s -> State s s
set = update . const

(<|>) :: State s a -> State s a -> State s a
ST a <|> ST b = ST (\s -> a s ++ b s)


-- label leafs of a tree in (not at all!) any order

data Tree a = Leaf a | Branch (Tree a) (Tree a)
 deriving Show

number :: Tree a -> [Tree (a,Int)]
number = runState 0 . numberST

numberST :: Tree a -> State Int (Tree (a,Int))
numberST (Leaf x) = do
  n <- update (+1)
  return (Leaf (x,n))

numberST (Branch l r)
  =  do nl <- numberST l
        nr <- numberST r
        return (Branch nl nr)
 <|> do nr <- numberST r
        nl <- numberST l
        return (Branch nl nr)