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