module ParseTree where

import Data.List
import Control.Monad (zipWithM)

data Symbol a
    = Constructor String a
    | Function    String a
    | Variable    String a
    
instance Show (Symbol a) where
    show (Constructor c _) = c
    show (Function    f _) = f
    show (Variable    v _) = v

instance Eq (Symbol a) where
    (==) (Constructor s1 _) (Constructor s2 _) = s1 == s2
    (==) (Function s1 _)    (Function s2 _)    = s1 == s2
    (==) (Variable s1 _)    (Variable s2 _)    = s1 == s2
    (==) _ _                                   = False

data SymbolTree a = SNode (Symbol a) [SymbolTree a]

instance Eq (SymbolTree a) where
    (==) (SNode (Constructor c _) ss) (SNode (Constructor c' _) ss') = 
        c == c' && length ss == length ss' && all (\(x,y) -> x == y) (zip ss ss')
    (==) (SNode (Function f _) ss) (SNode (Function f' _) ss') =
        f == f' && length ss == length ss' && all (\(x,y) -> x == y) (zip ss ss')
    (==) (SNode (Variable v _) _) (SNode (Variable v' _) _) = 
        v == v'

variables (SNode (Constructor c _) ss) = concat (map variables ss)
variables (SNode (Function f _)    ss) = concat (map variables ss)
variables (SNode s ss)                 = s:concat (map variables ss)

constructors s = nub $ cntrs s 
    where cntrs (SNode (s@(Constructor c _)) ss) = (s, length ss):concat (map cntrs ss)
          cntrs (SNode (Function f _) ss)        = concat (map cntrs ss)
          cntrs (SNode (Variable v _) _)         = []

data Declaration a 
    = Decl (SymbolTree a, SymbolTree a, [(SymbolTree a, SymbolTree a)]) Int -- LHS, RHS, WHERE CRAUSE

patternMatch :: SymbolTree a -> SymbolTree a -> Maybe [(String, SymbolTree a)]
patternMatch pattern value = 
    case (pattern, value) of 
      ( SNode (Constructor c a) es, SNode (Constructor c' a') es' ) 
          | c == c' ->
              do binds <- zipWithM patternMatch es es' 
                 return $ concat binds
      ( SNode (Function f a) es, SNode (Function f' a') es' ) 
          | f == f' ->
              do binds <- zipWithM patternMatch es es'
                 return $ concat binds
      ( SNode (Variable v _) _, t ) ->
          return $ [ (v, t) ]
      _ ->
          Nothing


substitute :: SymbolTree a -> [(String, SymbolTree a)] -> SymbolTree a 
substitute exp bind = 
    case exp of
      SNode (Constructor c a) es ->
          SNode (Constructor c a) $ map (\e -> substitute e bind) es
      SNode (Function f a) es ->
          SNode (Function f a) $ map (\e -> substitute e bind) es
      SNode (Variable v a) es ->
          case lookup v bind of 
            Just t   -> t
            Nothing  -> SNode (Variable v a) es

-- Utility Functions
lhs (Decl (l,r,_) _) = l
rhs (Decl (l,r,_) _) = r
wclause (Decl (_,_,w) _) = w
number (Decl _ n)    = n

defSymbol :: Declaration a -> Symbol a
defSymbol decl = let SNode f _ = lhs decl in f

defSymbols :: [Declaration a] -> [Symbol a]
defSymbols = nub.map defSymbol 

allConstructors :: [Declaration a] -> [(Symbol a, Int)]
allConstructors dcls = nub $ concat (map (constructors.rhs) dcls) ++ concat (map (constructors.lhs) dcls)


calledSymbols :: Declaration a -> [Symbol a] 
calledSymbols decl = cS (rhs decl)
    where cS (SNode (f@(Function _ _)) ts) = [f]
          cS (SNode _ ts)                  = concat $ map cS ts

calledFuncs :: Declaration a -> [SymbolTree a]
calledFuncs decl = cs (rhs decl)
    where cs (n@(SNode (Function _ _) ts)) = [n]
          cs (SNode _ ts)                  = concat $ map cs ts

bOf :: Symbol a -> [Declaration a]->[(Declaration a, Int)]
bOf s decls = filter (\(d,id) -> defSymbol d == s) $ zip decls [0..]

branchesOf :: Symbol a -> [Declaration a] -> [Declaration a]
branchesOf s decls = map fst $ bOf s decls

branchIDsOf :: Symbol a -> [Declaration a] -> [Int]
branchIDsOf s decls = map snd $ bOf s decls

 

leftVars :: Declaration a -> [Symbol a]
leftVars decl  = nub $ variables (lhs decl)

rightVars :: Declaration a -> [Symbol a]
rightVars decl = nub $ variables (rhs decl)

lostVars :: Declaration a -> [Symbol a]
lostVars decl = (leftVars decl) \\ (rightVars decl)



type ParseTree a = [Declaration a]

implode gr xs = concat$intersperse gr xs

instance Show (Declaration a) where
    show (Decl (l,r,w) b) = show l ++ " = " ++ show r ++
                            if length w > 0 then 
                                "where " ++ implode "\n      " (map show w) ++ "\n"
                            else 
                                "\n"
    showList xs = \s -> shows 0 xs++s
        where shows n xs = let s = replicate n ' ' in concat [ s++showN x | x <- xs] 
              n = maximum $ map (\(Decl (l,_,_) _)->length$show l) xs
              showN (Decl (l,r,w) b) = 
                  sN l ++ " = " ++ show r ++
                     if length w > 0 then
                         "\n  where \n" ++ showw 4 w
                     else 
                         "\n"
              sN l = let s = show l in s ++ replicate (n-(length s)) ' '
              showw nn xs = let s = replicate nn ' ' in concat [ s++showNW x | x <- xs ]
                  where n = maximum $ map (\(l,r)->length$show l) xs
                        showNW (l,r) = let s = show l in s ++ replicate (n-(length s)) ' ' ++ " = " ++ show r ++ "\n"
                        

instance Show (SymbolTree a) where
    show (SNode (Constructor c _) []) = c
    show (SNode (Constructor c _) cs) = c ++ "(" ++ (implode "," (map show cs)) ++ ")"
    show (SNode (Function f _) cs)    = f ++ "(" ++ (implode "," (map show cs)) ++ ")"
    show (SNode (Variable v _) _)     = v
