--
-- Setting up the state monad
--
data SM s a = Sm (s -> (a,s))

--
-- Runs the state monad on a state
--
runSM:: (SM s a) -> s -> (a,s)
runSM (Sm h) s = h s

instance Monad (SM s) where
     return a = Sm (\s -> (a,s))
     (Sm f) >>= g = Sm (\s -> (\(a,s') -> runSM (g a) s') (f s))
--
-- Examples:
--        using the expression data type
--

data Exp a b = Var b
             | Opn a [Exp a b]
     deriving (Eq,Show)

--
-- Numbering the leaves of an expression
--

number_exp :: (Exp a b) -> (SM Int (Exp a Int))
number_exp (Var x) 
   = do n <- Sm (\n -> (n,n+1))
        return (Var n)
number_exp (Opn f args) 
   = do args' <- number_args args
        return (Opn f args')

number_args :: [(Exp a b)] -> (SM Int [(Exp a Int)])
number_args [] = return []
number_args (x:xs) 
   = do x' <- number_exp x
        xs' <- number_args xs
        return (x':xs')
                   
ex1 = runSM (number_exp (Opn "f" [Var "a",Var "b",Var"a"])) 0
                
--
-- Substituting the variables of an expression with (new) numbers
--

type ISubs v = [(v,Int)]

data SF a = SS a | FF -- success or fail data type

--
--  the substitution of a string either replaces the string with a number as specified in 
--  the list of substitutions or adds a substitution using the least still number available ...
--
subst:: Eq v => v -> (SM (ISubs v) Int)
subst str = 
   Sm (\subs -> case (insubs str subs) of
                   SS n -> (n,subs)
                   FF ->(next_free str 0 subs))  where
         insubs:: Eq v => v -> (ISubs v) -> (SF Int)
         insubs str [] = FF
         insubs str ((str',n):subs)| str==str' = SS n
                                   | otherwise = insubs str subs
         next_free:: Eq v => v -> Int -> (ISubs v) -> (Int,ISubs v)
         next_free str n [] = (n,[(str,n)])
         next_free str n ((str',m):subs)
                    | n<m = (n,(str,n):((str',m):subs))
                    | otherwise = (\(m',subs') -> (m',((str',m):subs')))
                                  (next_free str (n+1) subs)

--                   
--  The code for substitution of an expression using the state monad ...
--
subst_exp:: Eq v => (Exp a v) -> (SM (ISubs v) (Exp a Int))
subst_exp (Var x) = do n <- subst x
                       return (Var n)
subst_exp (Opn f ts) = do ts' <- subst_list ts
                          return (Opn f ts')  where
         subst_list [] = return []
         subst_list (t:ts) = do t' <- subst_exp t
                                ts' <- subst_list ts
                                return (t':ts')

ex2 = runSM (subst_exp (Opn "f" [Var "a",Var "b",Var "a"])) []