on 727d708 help sign in
elk/compiler/Transform/Replace.hs
60 lines 2.1 kB history raw
1module Replace (replace, depth) where
3import Ast
4import Control.Arrow ((***))
5import Data.Set (singleton,empty,unions,member, Set)
7replace :: String -> Expr -> Expr -> Expr
8replace y v expr =
9 let f = replace y v in
10 case expr of
11 Range e1 e2 -> Range (f e1) (f e2)
12 Access e x -> Access (f e) x
13 Binop op e1 e2 -> Binop op (f e1) (f e2)
14 App e1 e2 -> App (f e1) (f e2)
15 If e1 e2 e3 -> If (f e1) (f e2) (f e3)
16 Guard ps -> map (f *** f) ps
17 Lift e es -> Lift (f e) (map f es)
18 Fold e1 e2 e3 -> Fold (f e1) (f e2) (f e3)
19 Async e -> Async (f e)
20 Let defs e -> if y `elem` vs then Let defs e else Let (zip vs (map f es)) (f e)
21 where (vs,es) = unzip defs
22 Var x -> if x == y then v else Var x
23 ExplicitList es -> ExplicitList (map f es)
24 Data name es -> Data name (map f es)
25 Case e cases -> Case (f e) $ map (caseReplace y v) cases
26 _ -> expr
28caseReplace :: String -> Expr -> (Pattern, Expr) -> (Pattern, Expr)
29caseReplace y v (p,e) =
30 if member y (patternVars p) then (p,e) else (p, replace y v e)
32patternVars :: Pattern -> Set String
33patternVars pattern =
34 case pattern of
35 PData _ ps -> unions (map patternVars ps)
36 PVar x -> singleton x
37 PAsVar x p -> insert x (patternVars p)
38 PAnything -> empty
39 PLiteral _ -> empty
41depth :: Expr -> Integer
42depth = depth' 0
43depth' d expr =
44 let f = depth' (d+1) in
45 case expr of
46 Range e1 e2 -> max (f e1) (f e2)
47 Access e x -> f e
48 Binop op e1 e2 -> max (f e1) (f e2)
49 Lambda x e -> f e
50 App e1 e2 -> max (f e1) (f e2)
51 If e1 e2 e3 -> maximum [f e1, f e2, f e3]
52 Guard ps -> maximum (map (uncurry max . f *** f) ps)
53 Lift e es -> maximum $ f e : map f es
54 Fold e1 e2 e3 -> maximum [f e1, f e2, f e3]
55 Async e -> f e
56 Let defs e -> let (_,es) = unzip defs in maximum $ f e : map f es
57 Data "::" es -> maximum $ map (depth' d) es
58 Data name es -> maximum $ 1 : map f es
59 Case e cases -> maximum $ f e : map (f . snd) cases
60 _ -> d