git.iko.soy/elk
on e411990 help sign in
elk/compiler/Transform/SortDefinitions.hs
198 lines 6.5 kB history raw
2module Transform.SortDefinitions (sortDefs, boundVars, flattenLets) where
4import Control.Monad.State
5import Control.Applicative ((<$>))
6import qualified Data.Set as Set
7import qualified Data.Map as Map
8import qualified SourceSyntax.Type as ST
9import SourceSyntax.Everything
10import qualified Data.Graph as Graph
11import qualified Data.Map as Map
12import qualified Data.Maybe as Maybe
14boundVars :: Pattern -> Set.Set String
15boundVars pattern =
16 case pattern of
17 PVar x -> Set.singleton x
18 PAlias x p -> Set.insert x (boundVars p)
19 PData _ ps -> Set.unions (map boundVars ps)
20 PRecord fields -> Set.fromList fields
21 PAnything -> Set.empty
22 PLiteral _ -> Set.empty
24ctors :: Pattern -> [String]
25ctors pattern =
26 case pattern of
27 PVar x -> []
28 PAlias x p -> ctors p
29 PData ctor ps -> ctor : concatMap ctors ps
30 PRecord fields -> []
31 PAnything -> []
32 PLiteral _ -> []
34free :: String -> State (Set.Set String) ()
35free x = modify (Set.insert x)
37bound :: Set.Set String -> State (Set.Set String) ()
38bound boundVars = modify (\freeVars -> Set.difference freeVars boundVars)
40sortDefs :: LExpr t v -> LExpr t v
41sortDefs expr = evalState (reorder expr) Set.empty
43flattenLets defs lexpr@(L _ expr) =
44 case expr of
45 Let ds body -> flattenLets (defs ++ ds) body
46 _ -> (defs, lexpr)
49reorder :: LExpr t v -> State (Set.Set String) (LExpr t v)
50reorder lexpr@(L s expr) =
51 L s `liftM`
52 case expr of
53 -- Be careful adding and restricting freeVars
54 Var x -> free x >> return expr
56 Lambda p e ->
57 uncurry Lambda `liftM` bindingReorder (p,e)
59 Binop op e1 e2 ->
60 do free op
61 Binop op `liftM` reorder e1 `ap` reorder e2
63 Case e cases ->
64 Case `liftM` reorder e `ap` mapM bindingReorder cases
66 Data name es ->
67 do free name
68 Data name `liftM` mapM reorder es
70 -- Just pipe the reorder though
71 Literal _ -> return expr
73 Range e1 e2 ->
74 Range `liftM` reorder e1 `ap` reorder e2
76 ExplicitList es ->
77 ExplicitList `liftM` mapM reorder es
79 App e1 e2 ->
80 App `liftM` reorder e1 `ap` reorder e2
82 MultiIf branches ->
83 MultiIf `liftM` mapM reorderPair branches
85 Access e lbl ->
86 Access `liftM` reorder e `ap` return lbl
88 Remove e lbl ->
89 Remove `liftM` reorder e `ap` return lbl
91 Insert e lbl v ->
92 Insert `liftM` reorder e `ap` return lbl `ap` reorder v
94 Modify e fields ->
95 Modify `liftM` reorder e `ap` mapM reorderField fields
97 Record fields ->
98 Record `liftM` mapM reorderField fields
100 Markdown uid md es -> Markdown uid md <$> mapM reorder es
102 -- Actually do some reordering
103 Let defs body ->
104 do body' <- reorder body
106 -- Sort defs into strongly connected components.This
107 -- allows the programmer to write definitions in whatever
108 -- order they please, we can still define things in order
109 -- and generalize polymorphic functions when appropriate.
110 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
111 let defss = map Graph.flattenSCC sccs
113 -- remove let-bound variables from the context
114 let getPatterns def =
115 case def of
116 Def pattern _ -> pattern
117 TypeAnnotation name _ -> PVar name
118 forM (map getPatterns defs) $ \pattern -> do
119 bound (boundVars pattern)
120 mapM free (ctors pattern)
122 let addDefs ds bod = L s (Let (concatMap toDefs ds) bod)
123 where
124 toDefs (pattern, expr, Nothing) = [ Def pattern expr ]
125 toDefs (PVar name, expr, Just tipe) =
126 [ TypeAnnotation name tipe, Def (PVar name) expr ]
128 L _ let' = foldr addDefs body' defss
130 return let'
133reorderField (label, expr) =
134 (,) label `liftM` reorder expr
136reorderPair (e1,e2) =
137 (,) `liftM` reorder e1 `ap` reorder e2
139bindingReorder :: (Pattern, LExpr t v) -> State (Set.Set String) (Pattern, LExpr t v)
140bindingReorder (pattern,expr) =
141 do expr' <- reorder expr
142 bound (boundVars pattern)
143 mapM free (ctors pattern)
144 return (pattern, expr')
147type PDef t v = (Pattern, LExpr t v, Maybe ST.Type)
149reorderAndGetDependencies :: PDef t v -> State (Set.Set String) (PDef t v, [String])
150reorderAndGetDependencies (pattern, expr, mType) =
151 do globalFrees <- get
152 -- work in a fresh environment
153 put Set.empty
154 expr' <- reorder expr
155 localFrees <- get
156 -- merge with global frees
157 modify (Set.union globalFrees)
158 return ((pattern, expr', mType), Set.toList localFrees)
161-- This also reorders the all of the sub-expressions in the Def list.
162buildDefDict :: [Def t v] -> State (Set.Set String) [(PDef t v, Int, [Int])]
163buildDefDict defs =
164 do pdefsDeps <- mapM reorderAndGetDependencies (getPDefs defs)
165 return $ realDeps (addKey pdefsDeps)
167 where
168 getPDefs :: [Def t v] -> [PDef t v]
169 getPDefs defs = map (\(p,(e,t)) -> (p,e,t)) $
170 Map.toList $ go defs Map.empty Map.empty
171 where
172 go [] ds ts =
173 Map.unions [ Map.difference ds ts
174 , Map.intersectionWith (\(e,_) t -> (e,Just t)) ds ts ]
176 go (def:defs) ds ts =
177 case def of
178 Def p e -> go defs (Map.insert p (e, Nothing) ds) ts
179 TypeAnnotation name tipe -> go defs ds (Map.insert (PVar name) tipe ts)
181 addKey :: [(PDef t v, [String])] -> [(PDef t v, Int, [String])]
182 addKey = zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
184 variableToKey :: (PDef t v, Int, [String]) -> [(String, Int)]
185 variableToKey ((pattern, _, _), key, _) =
186 [ (var, key) | var <- Set.toList (boundVars pattern) ]
188 variableToKeyMap :: [(PDef t v, Int, [String])] -> Map.Map String Int
189 variableToKeyMap pdefsDeps =
190 Map.fromList (concatMap variableToKey pdefsDeps)
192 realDeps :: [(PDef t v, Int, [String])] -> [(PDef t v, Int, [Int])]
193 realDeps pdefsDeps = map convert pdefsDeps
194 where
195 varDict = variableToKeyMap pdefsDeps
196 convert (pdef, key, deps) =
197 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)