git.iko.soy/elk
on 267f708 help sign in
elk/src/Canonicalize/Sort.hs
232 lines 5.9 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 ExplicitList es ->
72 ExplicitList <$> mapM reorder es
74 App func arg ->
75 App <$> reorder func <*> reorder arg
77 If branches finally ->
78 If
79 <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
80 <*> reorder finally
82 Access record field ->
83 Access
84 <$> reorder record
85 <*> return field
87 Update record fields ->
88 Update
89 <$> reorder record
90 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
92 Record fields ->
93 Record
94 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
96 Cmd _ ->
97 return expression
99 Sub _ ->
100 return expression
102 OutgoingPort _ _ ->
103 return expression
105 IncomingPort _ _ ->
106 return expression
108 Program _ _ ->
109 error "DANGER - Program AST nodes should not be in def sorting."
111 SaveEnv _ _ ->
112 return expression
114 GLShader _ _ _ ->
115 return expression
117 -- Actually do some reordering
118 Let defs body ->
119 do body' <- reorder body
121 -- Sort defs into strongly connected components.This
122 -- allows the programmer to write definitions in whatever
123 -- order they please, we can still define things in order
124 -- and generalize polymorphic functions when appropriate.
125 sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
126 let defss = map Graph.flattenSCC sccs
128 -- remove let-bound variables from the context
129 forM_ defs $ \(Canonical.Def _ pattern _ _) -> do
130 bound pattern
131 mapM free (ctors pattern)
133 let (A.A _ let') =
134 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
136 return let'
139ctors :: P.Canonical -> [String]
140ctors (A.A _ pattern) =
141 case pattern of
142 P.Var _ ->
143 []
145 P.Alias _ p ->
146 ctors p
148 P.Record _ ->
149 []
151 P.Anything ->
152 []
154 P.Literal _ ->
155 []
157 P.Data (V.Canonical home name) ps ->
158 case home of
159 V.Local -> name : rest
160 V.TopLevel _ -> name : rest
161 V.BuiltIn -> rest
162 V.Module _ -> rest
163 where
164 rest = concatMap ctors ps
167bound :: P.Canonical -> State (Set.Set String) ()
168bound pattern =
169 let
170 boundVars =
171 P.boundVarSet pattern
172 in
173 modify (\freeVars -> Set.difference freeVars boundVars)
176bindingReorder
177 :: (P.Canonical, Canonical.Expr)
178 -> State (Set.Set String) (P.Canonical, Canonical.Expr)
179bindingReorder (pattern, expr) =
180 do expr' <- reorder expr
181 bound pattern
182 mapM_ free (ctors pattern)
183 return (pattern, expr')
187-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS
190-- This also reorders the all of the sub-expressions in the Def list.
191buildDefGraph
192 :: [Canonical.Def]
193 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
194buildDefGraph defs =
195 do pdefsDeps <- mapM reorderAndGetDependencies defs
196 return $ realDeps (addKey pdefsDeps)
197 where
198 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
199 addKey =
200 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
202 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
203 variableToKey (Canonical.Def _ pattern _ _, key, _) =
204 [ (var, key) | var <- P.boundVarList pattern ]
206 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
207 variableToKeyMap pdefsDeps =
208 Map.fromList (concatMap variableToKey pdefsDeps)
210 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
211 realDeps pdefsDeps =
212 map convert pdefsDeps
213 where
214 varDict =
215 variableToKeyMap pdefsDeps
217 convert (pdef, key, deps) =
218 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)
221reorderAndGetDependencies
222 :: Canonical.Def
223 -> State (Set.Set String) (Canonical.Def, [String])
224reorderAndGetDependencies (Canonical.Def facts pattern expr mType) =
225 do globalFrees <- get
226 -- work in a fresh environment
227 put Set.empty
228 expr' <- reorder expr
229 localFrees <- get
230 -- merge with global frees
231 modify (Set.union globalFrees)
232 return (Canonical.Def facts pattern expr' mType, Set.toList localFrees)