git.iko.soy/elk
on 8913c30 help sign in
elk/src/Transform/SortDefinitions.hs
193 lines 5.9 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 Port name impl ->
119 return $ Port name impl
121 -- Actually do some reordering
122 Let defs body ->
123 do body' <- reorder body
125 -- Sort defs into strongly connected components.This
126 -- allows the programmer to write definitions in whatever
127 -- order they please, we can still define things in order
128 -- and generalize polymorphic functions when appropriate.
129 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
130 let defss = map Graph.flattenSCC sccs
132 -- remove let-bound variables from the context
133 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
134 bound (P.boundVars pattern)
135 mapM free (ctors pattern)
137 let A _ let' =
138 foldr (\ds bod -> A ann (Let ds bod)) body' defss
140 return let'
143bindingReorder
144 :: (P.CanonicalPattern, Canonical.Expr)
145 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
146bindingReorder (pattern,expr) =
147 do expr' <- reorder expr
148 bound (P.boundVars pattern)
149 mapM_ free (ctors pattern)
150 return (pattern, expr')
153reorderAndGetDependencies
154 :: Canonical.Def
155 -> State (Set.Set String) (Canonical.Def, [String])
156reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
157 do globalFrees <- get
158 -- work in a fresh environment
159 put Set.empty
160 expr' <- reorder expr
161 localFrees <- get
162 -- merge with global frees
163 modify (Set.union globalFrees)
164 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
167-- This also reorders the all of the sub-expressions in the Def list.
168buildDefDict
169 :: [Canonical.Def]
170 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
171buildDefDict defs =
172 do pdefsDeps <- mapM reorderAndGetDependencies defs
173 return $ realDeps (addKey pdefsDeps)
174 where
175 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
176 addKey =
177 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
179 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
180 variableToKey (Canonical.Definition pattern _ _, key, _) =
181 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
183 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
184 variableToKeyMap pdefsDeps =
185 Map.fromList (concatMap variableToKey pdefsDeps)
187 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
188 realDeps pdefsDeps =
189 map convert pdefsDeps
190 where
191 varDict = variableToKeyMap pdefsDeps
192 convert (pdef, key, deps) =
193 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)