git.iko.soy/elk
on 294f215 help sign in
elk/src/Transform/SortDefinitions.hs
200 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 Loopback name tipe maybeExpr ->
125 Loopback name tipe <$>
126 maybe (return Nothing) (fmap Just . reorder) maybeExpr
128 -- Actually do some reordering
129 Let defs body ->
130 do body' <- reorder body
132 -- Sort defs into strongly connected components.This
133 -- allows the programmer to write definitions in whatever
134 -- order they please, we can still define things in order
135 -- and generalize polymorphic functions when appropriate.
136 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
137 let defss = map Graph.flattenSCC sccs
139 -- remove let-bound variables from the context
140 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
141 bound (P.boundVars pattern)
142 mapM free (ctors pattern)
144 let A _ let' =
145 foldr (\ds bod -> A ann (Let ds bod)) body' defss
147 return let'
150bindingReorder
151 :: (P.CanonicalPattern, Canonical.Expr)
152 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
153bindingReorder (pattern,expr) =
154 do expr' <- reorder expr
155 bound (P.boundVars pattern)
156 mapM_ free (ctors pattern)
157 return (pattern, expr')
160reorderAndGetDependencies
161 :: Canonical.Def
162 -> State (Set.Set String) (Canonical.Def, [String])
163reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
164 do globalFrees <- get
165 -- work in a fresh environment
166 put Set.empty
167 expr' <- reorder expr
168 localFrees <- get
169 -- merge with global frees
170 modify (Set.union globalFrees)
171 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
174-- This also reorders the all of the sub-expressions in the Def list.
175buildDefDict
176 :: [Canonical.Def]
177 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
178buildDefDict defs =
179 do pdefsDeps <- mapM reorderAndGetDependencies defs
180 return $ realDeps (addKey pdefsDeps)
181 where
182 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
183 addKey =
184 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
186 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
187 variableToKey (Canonical.Definition pattern _ _, key, _) =
188 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
190 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
191 variableToKeyMap pdefsDeps =
192 Map.fromList (concatMap variableToKey pdefsDeps)
194 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
195 realDeps pdefsDeps =
196 map convert pdefsDeps
197 where
198 varDict = variableToKeyMap pdefsDeps
199 convert (pdef, key, deps) =
200 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)