git.iko.soy/elk
on da346ea help sign in
elk/compiler/Transform/SortDefinitions.hs
185 lines 6.1 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
24free :: String -> State (Set.Set String) ()
25free x = modify (Set.insert x)
27bound :: Set.Set String -> State (Set.Set String) ()
28bound boundVars = modify (\freeVars -> Set.difference freeVars boundVars)
30sortDefs :: LExpr t v -> LExpr t v
31sortDefs expr = evalState (reorder expr) Set.empty
33flattenLets defs lexpr@(L _ _ expr) =
34 case expr of
35 Let ds body -> flattenLets (defs ++ ds) body
36 _ -> (defs, lexpr)
39reorder :: LExpr t v -> State (Set.Set String) (LExpr t v)
40reorder lexpr@(L a b expr) =
41 L a b `liftM`
42 case expr of
43 -- Be careful adding and restricting freeVars
44 Var x -> free x >> return expr
46 Lambda p e ->
47 uncurry Lambda `liftM` bindingReorder (p,e)
49 Binop op e1 e2 ->
50 do free op
51 Binop op `liftM` reorder e1 `ap` reorder e2
53 Case e cases ->
54 Case `liftM` reorder e `ap` mapM bindingReorder cases
56 Data name es ->
57 do free name
58 Data name `liftM` mapM reorder es
60 -- Just pipe the reorder though
61 Literal _ -> return expr
63 Range e1 e2 ->
64 Range `liftM` reorder e1 `ap` reorder e2
66 ExplicitList es ->
67 ExplicitList `liftM` mapM reorder es
69 App e1 e2 ->
70 App `liftM` reorder e1 `ap` reorder e2
72 MultiIf branches ->
73 MultiIf `liftM` mapM reorderPair branches
75 Access e lbl ->
76 Access `liftM` reorder e `ap` return lbl
78 Remove e lbl ->
79 Remove `liftM` reorder e `ap` return lbl
81 Insert e lbl v ->
82 Insert `liftM` reorder e `ap` return lbl `ap` reorder v
84 Modify e fields ->
85 Modify `liftM` reorder e `ap` mapM reorderField fields
87 Record fields ->
88 Record `liftM` mapM reorderField fields
90 Markdown _ -> return expr
92 -- Actually do some reordering
93 Let defs body ->
94 do body' <- reorder body
96 -- Sort defs into strongly connected components.This
97 -- allows the programmer to write definitions in whatever
98 -- order they please, we can still define things in order
99 -- and generalize polymorphic functions when appropriate.
100 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
101 let defss = map Graph.flattenSCC sccs
103 -- remove let-bound variables from the context
104 let getPatterns def =
105 case def of
106 Def pattern _ -> pattern
107 TypeAnnotation name _ -> PVar name
108 mapM (bound . boundVars . getPatterns) defs
110 let addDefs ds bod = L a b (Let (concatMap toDefs ds) bod)
111 where
112 toDefs (pattern, expr, Nothing) = [ Def pattern expr ]
113 toDefs (PVar name, expr, Just tipe) =
114 [ TypeAnnotation name tipe, Def (PVar name) expr ]
116 L _ _ let' = foldr addDefs body' defss
118 return let'
121reorderField (label, expr) =
122 (,) label `liftM` reorder expr
124reorderPair (e1,e2) =
125 (,) `liftM` reorder e1 `ap` reorder e2
127bindingReorder :: (Pattern, LExpr t v) -> State (Set.Set String) (Pattern, LExpr t v)
128bindingReorder (pattern,expr) =
129 do expr' <- reorder expr
130 bound (boundVars pattern)
131 return (pattern, expr')
134type PDef t v = (Pattern, LExpr t v, Maybe ST.Type)
136reorderAndGetDependencies :: PDef t v -> State (Set.Set String) (PDef t v, [String])
137reorderAndGetDependencies (pattern, expr, mType) =
138 do globalFrees <- get
139 -- work in a fresh environment
140 put Set.empty
141 expr' <- reorder expr
142 localFrees <- get
143 -- merge with global frees
144 modify (Set.union globalFrees)
145 return ((pattern, expr', mType), Set.toList localFrees)
148-- This also reorders the all of the sub-expressions in the Def list.
149buildDefDict :: [Def t v] -> State (Set.Set String) [(PDef t v, Int, [Int])]
150buildDefDict defs =
151 do pdefsDeps <- mapM reorderAndGetDependencies (getPDefs defs)
152 return $ realDeps (addKey pdefsDeps)
154 where
155 getPDefs :: [Def t v] -> [PDef t v]
156 getPDefs defs = map (\(p,(e,t)) -> (p,e,t)) $
157 Map.toList $ go defs Map.empty Map.empty
158 where
159 go [] ds ts =
160 Map.unions [ Map.difference ds ts
161 , Map.intersectionWith (\(e,_) t -> (e,Just t)) ds ts ]
163 go (def:defs) ds ts =
164 case def of
165 Def p e -> go defs (Map.insert p (e, Nothing) ds) ts
166 TypeAnnotation name tipe -> go defs ds (Map.insert (PVar name) tipe ts)
168 addKey :: [(PDef t v, [String])] -> [(PDef t v, Int, [String])]
169 addKey = zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
171 variableToKey :: (PDef t v, Int, [String]) -> [(String, Int)]
172 variableToKey ((pattern, _, _), key, _) =
173 [ (var, key) | var <- Set.toList (boundVars pattern) ]
175 variableToKeyMap :: [(PDef t v, Int, [String])] -> Map.Map String Int
176 variableToKeyMap pdefsDeps =
177 Map.fromList (concatMap variableToKey pdefsDeps)
179 realDeps :: [(PDef t v, Int, [String])] -> [(PDef t v, Int, [Int])]
180 realDeps pdefsDeps = map convert pdefsDeps
181 where
182 varDict = variableToKeyMap pdefsDeps
183 convert (pdef, key, deps) =
184 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)