git.iko.soy/elk
on 357dedf help sign in
elk/src/Transform/SortDefinitions.hs
214 lines 6.2 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Transform.SortDefinitions (sortDefs) where
4import Control.Monad.State
5import Control.Applicative ((<$>),(<*>))
6import qualified Data.Graph as Graph
7import qualified Data.Map as Map
8import qualified Data.Maybe as Maybe
9import qualified Data.Set as Set
11import AST.Expression.General (Expr'(..), PortImpl(..))
12import qualified AST.Expression.Canonical as Canonical
13import qualified AST.Pattern as P
14import qualified AST.Variable as V
15import qualified Reporting.Annotation as A
18ctors :: P.CanonicalPattern -> [String]
19ctors (A.A _ pattern) =
20 case pattern of
21 P.Var _ ->
22 []
24 P.Alias _ p ->
25 ctors p
27 P.Record _ ->
28 []
30 P.Anything ->
31 []
33 P.Literal _ ->
34 []
36 P.Data (V.Canonical home name) ps ->
37 case home of
38 V.Local -> name : rest
39 V.BuiltIn -> rest
40 V.Module _ -> rest
41 where
42 rest = concatMap ctors ps
45free :: String -> State (Set.Set String) ()
46free x =
47 modify (Set.insert x)
50freeIfLocal :: V.Canonical -> State (Set.Set String) ()
51freeIfLocal (V.Canonical home name) =
52 case home of
53 V.Local -> free name
54 V.BuiltIn -> return ()
55 V.Module _ -> return ()
58bound :: P.CanonicalPattern -> State (Set.Set String) ()
59bound pattern =
60 let boundVars = P.boundVarSet pattern
61 in
62 modify (\freeVars -> Set.difference freeVars boundVars)
65sortDefs :: Canonical.Expr -> Canonical.Expr
66sortDefs expression =
67 evalState (reorder expression) Set.empty
70reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
71reorder (A.A ann expression) =
72 A.A ann <$>
73 case expression of
74 -- Be careful adding and restricting freeVars
75 Var var ->
76 do freeIfLocal var
77 return expression
79 Lambda pattern body ->
80 uncurry Lambda <$> bindingReorder (pattern,body)
82 Binop op leftExpr rightExpr ->
83 do freeIfLocal op
84 Binop op <$> reorder leftExpr <*> reorder rightExpr
86 Case expr cases ->
87 Case <$> reorder expr <*> mapM bindingReorder cases
89 Data name exprs ->
90 do free name
91 Data name <$> mapM reorder exprs
93 -- Just pipe the reorder though
94 Literal _ ->
95 return expression
97 Range lowExpr highExpr ->
98 Range <$> reorder lowExpr <*> reorder highExpr
100 ExplicitList es ->
101 ExplicitList <$> mapM reorder es
103 App func arg ->
104 App <$> reorder func <*> reorder arg
106 MultiIf branches ->
107 MultiIf <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
109 Access record field ->
110 Access <$> reorder record <*> return field
112 Remove record field ->
113 Remove <$> reorder record <*> return field
115 Insert record field expr ->
116 Insert <$> reorder record <*> return field <*> reorder expr
118 Modify record fields ->
119 Modify
120 <$> reorder record
121 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
123 Record fields ->
124 Record
125 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
127 GLShader _ _ _ ->
128 return expression
130 Port impl ->
131 Port <$>
132 case impl of
133 In _ _ ->
134 return impl
136 Out name expr tipe ->
137 (\e -> Out name e tipe) <$> reorder expr
139 Task name expr tipe ->
140 (\e -> Task name e tipe) <$> reorder expr
142 -- Actually do some reordering
143 Let defs body ->
144 do body' <- reorder body
146 -- Sort defs into strongly connected components.This
147 -- allows the programmer to write definitions in whatever
148 -- order they please, we can still define things in order
149 -- and generalize polymorphic functions when appropriate.
150 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
151 let defss = map Graph.flattenSCC sccs
153 -- remove let-bound variables from the context
154 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
155 bound pattern
156 mapM free (ctors pattern)
158 let (A.A _ let') =
159 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
161 return let'
164bindingReorder
165 :: (P.CanonicalPattern, Canonical.Expr)
166 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
167bindingReorder (pattern,expr) =
168 do expr' <- reorder expr
169 bound pattern
170 mapM_ free (ctors pattern)
171 return (pattern, expr')
174reorderAndGetDependencies
175 :: Canonical.Def
176 -> State (Set.Set String) (Canonical.Def, [String])
177reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
178 do globalFrees <- get
179 -- work in a fresh environment
180 put Set.empty
181 expr' <- reorder expr
182 localFrees <- get
183 -- merge with global frees
184 modify (Set.union globalFrees)
185 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
188-- This also reorders the all of the sub-expressions in the Def list.
189buildDefDict
190 :: [Canonical.Def]
191 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
192buildDefDict defs =
193 do pdefsDeps <- mapM reorderAndGetDependencies defs
194 return $ realDeps (addKey pdefsDeps)
195 where
196 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
197 addKey =
198 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
200 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
201 variableToKey (Canonical.Definition pattern _ _, key, _) =
202 [ (var, key) | var <- P.boundVarList pattern ]
204 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
205 variableToKeyMap pdefsDeps =
206 Map.fromList (concatMap variableToKey pdefsDeps)
208 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
209 realDeps pdefsDeps =
210 map convert pdefsDeps
211 where
212 varDict = variableToKeyMap pdefsDeps
213 convert (pdef, key, deps) =
214 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)