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

Tree @master (Download .tar.gz)

AstUtils.hs @master — raw · history · blame 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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