git.iko.soy/elk
on fa8ade2 help sign in
elk/src/Canonicalize/Sort.hs
235 lines 6.0 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Canonicalize.Sort (definitions) where
4import Control.Monad.State
5import qualified Data.Graph as Graph
6import qualified Data.Map as Map
7import qualified Data.Maybe as Maybe
8import qualified Data.Set as Set
10import AST.Expression.General (Expr'(..))
11import qualified AST.Expression.Canonical as Canonical
12import qualified AST.Pattern as P
13import qualified AST.Variable as V
14import qualified Reporting.Annotation as A
18-- STARTING POINT
21definitions :: Canonical.Expr -> Canonical.Expr
22definitions expression =
23 evalState (reorder expression) Set.empty
26free :: String -> State (Set.Set String) ()
27free x =
28 modify (Set.insert x)
31freeIfLocal :: V.Canonical -> State (Set.Set String) ()
32freeIfLocal (V.Canonical home name) =
33 case home of
34 V.Local -> free name
35 V.TopLevel _ -> free name
36 V.BuiltIn -> return ()
37 V.Module _ -> return ()
41-- REORDER EXPRESSIONS
44reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
45reorder (A.A ann expression) =
46 A.A ann <$>
47 case expression of
48 -- Be careful adding and restricting freeVars
49 Var var ->
50 do freeIfLocal var
51 return expression
53 Lambda pattern body ->
54 uncurry Lambda <$> bindingReorder (pattern,body)
56 Binop op leftExpr rightExpr ->
57 do freeIfLocal op
58 Binop op <$> reorder leftExpr <*> reorder rightExpr
60 Case expr cases ->
61 Case <$> reorder expr <*> mapM bindingReorder cases
63 Data name exprs ->
64 do free name
65 Data name <$> mapM reorder exprs
67 -- Just pipe the reorder though
68 Literal _ ->
69 return expression
71 Range lowExpr highExpr ->
72 Range <$> reorder lowExpr <*> reorder highExpr
74 ExplicitList es ->
75 ExplicitList <$> mapM reorder es
77 App func arg ->
78 App <$> reorder func <*> reorder arg
80 If branches finally ->
81 If
82 <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
83 <*> reorder finally
85 Access record field ->
86 Access
87 <$> reorder record
88 <*> return field
90 Update record fields ->
91 Update
92 <$> reorder record
93 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
95 Record fields ->
96 Record
97 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
99 Cmd _ ->
100 return expression
102 Sub _ ->
103 return expression
105 OutgoingPort _ _ ->
106 return expression
108 IncomingPort _ _ ->
109 return expression
111 Program _ _ ->
112 error "DANGER - Program AST nodes should not be in def sorting."
114 SaveEnv _ _ ->
115 return expression
117 GLShader _ _ _ ->
118 return expression
120 -- Actually do some reordering
121 Let defs body ->
122 do body' <- reorder body
124 -- Sort defs into strongly connected components.This
125 -- allows the programmer to write definitions in whatever
126 -- order they please, we can still define things in order
127 -- and generalize polymorphic functions when appropriate.
128 sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
129 let defss = map Graph.flattenSCC sccs
131 -- remove let-bound variables from the context
132 forM_ defs $ \(Canonical.Def _ pattern _ _) -> do
133 bound pattern
134 mapM free (ctors pattern)
136 let (A.A _ let') =
137 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
139 return let'
142ctors :: P.Canonical -> [String]
143ctors (A.A _ pattern) =
144 case pattern of
145 P.Var _ ->
146 []
148 P.Alias _ p ->
149 ctors p
151 P.Record _ ->
152 []
154 P.Anything ->
155 []
157 P.Literal _ ->
158 []
160 P.Data (V.Canonical home name) ps ->
161 case home of
162 V.Local -> name : rest
163 V.TopLevel _ -> name : rest
164 V.BuiltIn -> rest
165 V.Module _ -> rest
166 where
167 rest = concatMap ctors ps
170bound :: P.Canonical -> State (Set.Set String) ()
171bound pattern =
172 let
173 boundVars =
174 P.boundVarSet pattern
175 in
176 modify (\freeVars -> Set.difference freeVars boundVars)
179bindingReorder
180 :: (P.Canonical, Canonical.Expr)
181 -> State (Set.Set String) (P.Canonical, Canonical.Expr)
182bindingReorder (pattern, expr) =
183 do expr' <- reorder expr
184 bound pattern
185 mapM_ free (ctors pattern)
186 return (pattern, expr')
190-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS
193-- This also reorders the all of the sub-expressions in the Def list.
194buildDefGraph
195 :: [Canonical.Def]
196 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
197buildDefGraph defs =
198 do pdefsDeps <- mapM reorderAndGetDependencies defs
199 return $ realDeps (addKey pdefsDeps)
200 where
201 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
202 addKey =
203 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
205 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
206 variableToKey (Canonical.Def _ pattern _ _, key, _) =
207 [ (var, key) | var <- P.boundVarList pattern ]
209 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
210 variableToKeyMap pdefsDeps =
211 Map.fromList (concatMap variableToKey pdefsDeps)
213 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
214 realDeps pdefsDeps =
215 map convert pdefsDeps
216 where
217 varDict =
218 variableToKeyMap pdefsDeps
220 convert (pdef, key, deps) =
221 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)
224reorderAndGetDependencies
225 :: Canonical.Def
226 -> State (Set.Set String) (Canonical.Def, [String])
227reorderAndGetDependencies (Canonical.Def facts pattern expr mType) =
228 do globalFrees <- get
229 -- work in a fresh environment
230 put Set.empty
231 expr' <- reorder expr
232 localFrees <- get
233 -- merge with global frees
234 modify (Set.union globalFrees)
235 return (Canonical.Def facts pattern expr' mType, Set.toList localFrees)