git.iko.soy/elk
on 552becc help sign in
elk/src/Transform/SortDefinitions.hs
202 lines 6.1 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.Annotation
12import AST.Expression.General (Expr'(..))
13import qualified AST.Expression.Canonical as Canonical
14import qualified AST.Pattern as P
15import qualified AST.Variable as V
18ctors :: P.CanonicalPattern -> [String]
19ctors pattern =
20 case pattern of
21 P.Var _ -> []
22 P.Alias _ p -> ctors p
23 P.Record _ -> []
24 P.Anything -> []
25 P.Literal _ -> []
26 P.Data (V.Canonical home name) ps ->
27 case home of
28 V.Local -> name : rest
29 V.BuiltIn -> rest
30 V.Module _ -> rest
31 where
32 rest = concatMap ctors ps
35free :: String -> State (Set.Set String) ()
36free x =
37 modify (Set.insert x)
40freeIfLocal :: V.Canonical -> State (Set.Set String) ()
41freeIfLocal (V.Canonical home name) =
42 case home of
43 V.Local -> free name
44 V.BuiltIn -> return ()
45 V.Module _ -> return ()
48bound :: Set.Set String -> State (Set.Set String) ()
49bound boundVars =
50 modify (\freeVars -> Set.difference freeVars boundVars)
53sortDefs :: Canonical.Expr -> Canonical.Expr
54sortDefs expression =
55 evalState (reorder expression) Set.empty
58reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
59reorder (A ann expression) =
60 A ann <$>
61 case expression of
62 -- Be careful adding and restricting freeVars
63 Var var ->
64 do freeIfLocal var
65 return expression
67 Lambda pattern body ->
68 uncurry Lambda <$> bindingReorder (pattern,body)
70 Binop op leftExpr rightExpr ->
71 do freeIfLocal op
72 Binop op <$> reorder leftExpr <*> reorder rightExpr
74 Case expr cases ->
75 Case <$> reorder expr <*> mapM bindingReorder cases
77 Data name exprs ->
78 do free name
79 Data name <$> mapM reorder exprs
81 -- Just pipe the reorder though
82 Literal _ ->
83 return expression
85 Range lowExpr highExpr ->
86 Range <$> reorder lowExpr <*> reorder highExpr
88 ExplicitList es ->
89 ExplicitList <$> mapM reorder es
91 App func arg ->
92 App <$> reorder func <*> reorder arg
94 MultiIf branches ->
95 MultiIf <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
97 Access record field ->
98 Access <$> reorder record <*> return field
100 Remove record field ->
101 Remove <$> reorder record <*> return field
103 Insert record field expr ->
104 Insert <$> reorder record <*> return field <*> reorder expr
106 Modify record fields ->
107 Modify
108 <$> reorder record
109 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
111 Record fields ->
112 Record
113 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
115 GLShader _ _ _ ->
116 return expression
118 Input name tipe ->
119 return $ Input name tipe
121 Output name tipe signal ->
122 Output name tipe <$> reorder signal
124 LoopbackIn name source ->
125 return $ LoopbackIn name source
127 LoopbackOut name expr ->
128 LoopbackOut name <$> reorder expr
130 -- Actually do some reordering
131 Let defs body ->
132 do body' <- reorder body
134 -- Sort defs into strongly connected components.This
135 -- allows the programmer to write definitions in whatever
136 -- order they please, we can still define things in order
137 -- and generalize polymorphic functions when appropriate.
138 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
139 let defss = map Graph.flattenSCC sccs
141 -- remove let-bound variables from the context
142 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
143 bound (P.boundVars pattern)
144 mapM free (ctors pattern)
146 let A _ let' =
147 foldr (\ds bod -> A ann (Let ds bod)) body' defss
149 return let'
152bindingReorder
153 :: (P.CanonicalPattern, Canonical.Expr)
154 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
155bindingReorder (pattern,expr) =
156 do expr' <- reorder expr
157 bound (P.boundVars pattern)
158 mapM_ free (ctors pattern)
159 return (pattern, expr')
162reorderAndGetDependencies
163 :: Canonical.Def
164 -> State (Set.Set String) (Canonical.Def, [String])
165reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
166 do globalFrees <- get
167 -- work in a fresh environment
168 put Set.empty
169 expr' <- reorder expr
170 localFrees <- get
171 -- merge with global frees
172 modify (Set.union globalFrees)
173 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
176-- This also reorders the all of the sub-expressions in the Def list.
177buildDefDict
178 :: [Canonical.Def]
179 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
180buildDefDict defs =
181 do pdefsDeps <- mapM reorderAndGetDependencies defs
182 return $ realDeps (addKey pdefsDeps)
183 where
184 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
185 addKey =
186 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
188 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
189 variableToKey (Canonical.Definition pattern _ _, key, _) =
190 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
192 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
193 variableToKeyMap pdefsDeps =
194 Map.fromList (concatMap variableToKey pdefsDeps)
196 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
197 realDeps pdefsDeps =
198 map convert pdefsDeps
199 where
200 varDict = variableToKeyMap pdefsDeps
201 convert (pdef, key, deps) =
202 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)