git.iko.soy/elk
on 94ebc2d 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 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 ForeignCmd _ _ ->
106 return expression
108 ForeignSub _ _ ->
109 return expression
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)