git.haldean.org del / master src / Language / Del / AstUtils.hs
master

Tree @master (Download .tar.gz)

AstUtils.hs @masterraw · history · blame

module Language.Del.AstUtils where

import           Data.Functor.Identity
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import           Language.Del.AST

replaceNames :: Map.Map T.Text Expr -> Expr -> Expr
replaceNames pmap e@(Id name) = Map.findWithDefault e name pmap
replaceNames _ e = e

swapNames :: Map.Map T.Text T.Text -> Expr -> Expr
swapNames m = replaceNames $ Map.map (\name -> Id name) m

containsName :: Name -> Expr -> Bool
containsName n = anyexpr (\e -> case e of
    Id name -> name == n
    _ -> False)

-- Calls the provided transformation function on each node in the AST.
mapexprM :: Monad m => (Expr -> m Expr) -> Expr -> m Expr
mapexprM f (Add exprs) = mapM (mapexprM f) exprs >>= f . Add
mapexprM f (Mul exprs) = mapM (mapexprM f) exprs >>= f . Mul
mapexprM f (Sub e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ Sub m1 m2
mapexprM f (Div e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ Div m1 m2
mapexprM f (Pow e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ Pow m1 m2
mapexprM f (Gt e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ Gt m1 m2
mapexprM f (Lt e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ Lt m1 m2
mapexprM f (GtE e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ GtE m1 m2
mapexprM f (LtE e1 e2) = do
    m1 <- mapexprM f e1
    m2 <- mapexprM f e2
    f $ LtE m1 m2
mapexprM f (Subscript e sub) = mapexprM f e >>= \x -> f $ Subscript x sub
mapexprM f (Negate e) = mapexprM f e >>= \x -> f $ Negate x
mapexprM f (Apply hd xs) = do
    mh <- mapexprM f hd
    mb <- mapM (mapexprM f) xs
    f $ Apply mh mb 
mapexprM f x@(Id _) = f x
mapexprM f x@(ScopedId _ _) = f x
mapexprM f (Number num) = f $ Number num
mapexprM f (Let n e b) = do
    me <- mapexprM f e
    mb <- mapexprM f b
    f $ Let n me mb
mapexprM f b@(Boolean _) = f b
mapexprM f (If c bt bf) = do
    mc <- mapexprM f c
    mt <- mapexprM f bt
    mf <- mapexprM f bf
    f $ If mc mt mf
mapexprM f (NativeCall n xs) = mapM (mapexprM f) xs >>= f . (NativeCall n)

mapexpr :: (Expr -> Expr) -> Expr -> Expr
mapexpr f = runIdentity . mapexprM (return . f)

mapast :: (Def -> Def) -> AST -> AST
mapast f (AST m i) = AST (Map.map f m) i

mapdefs :: (Expr -> Expr) -> AST -> AST
mapdefs f = mapast (\(Def n ret p b) -> Def n ret p (f b))

mapastM :: (Monad m) => (Def -> m Def) -> AST -> m AST
mapastM f (AST m i)
    = newpairs >>= return . (flip AST $ i) . Map.fromList
    where newpairs = mapM (\(k, v) -> f v >>= \newv -> return (k, newv)) $ Map.toList m

-- Returns true if the provided property function returns true for any
-- expression in the tree.
anyexpr :: (Expr -> Bool) -> Expr -> Bool
anyexpr f e@(Add exprs) = if f e then True else any id $ map (anyexpr f) exprs
anyexpr f e@(Mul exprs) = if f e then True else any id $ map (anyexpr f) exprs
anyexpr f e@(Sub e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(Div e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(Pow e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(Gt e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(Lt e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(GtE e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(LtE e1 e2) = if f e then True else anyexpr f e1 || anyexpr f e2
anyexpr f e@(Subscript ex _) = if f e then True else anyexpr f ex
anyexpr f e@(Negate ex) = if f e then True else anyexpr f ex
anyexpr f e@(Apply hd xs) = if f e then True else anyexpr f hd || any id (map (anyexpr f) xs)
anyexpr f e@(Id _) = f e
anyexpr f e@(ScopedId _ _) = f e
anyexpr f e@(Number _) = f e
anyexpr f e@(Let _ v b) = if f e then True else anyexpr f v || anyexpr f b
anyexpr f e@(Boolean _) = f e
anyexpr f e@(If c bt bf) = if f e then True else anyexpr f c || anyexpr f bt || anyexpr f bf
anyexpr f e@(NativeCall _ xs) = if f e then True else any id $ map (anyexpr f) xs