git.iko.soy/elk
on 07be0d2 help sign in
elk/src/Canonicalize/Sort.hs
223 lines 5.8 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 GLShader _ _ _ ->
106 return expression
108 -- Actually do some reordering
109 Let defs body ->
110 do body' <- reorder body
112 -- Sort defs into strongly connected components.This
113 -- allows the programmer to write definitions in whatever
114 -- order they please, we can still define things in order
115 -- and generalize polymorphic functions when appropriate.
116 sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
117 let defss = map Graph.flattenSCC sccs
119 -- remove let-bound variables from the context
120 forM_ defs $ \(Canonical.Definition _ pattern _ _) -> do
121 bound pattern
122 mapM free (ctors pattern)
124 let (A.A _ let') =
125 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
127 return let'
130ctors :: P.Canonical -> [String]
131ctors (A.A _ pattern) =
132 case pattern of
133 P.Var _ ->
134 []
136 P.Alias _ p ->
137 ctors p
139 P.Record _ ->
140 []
142 P.Anything ->
143 []
145 P.Literal _ ->
146 []
148 P.Data (V.Canonical home name) ps ->
149 case home of
150 V.Local -> name : rest
151 V.TopLevel _ -> name : rest
152 V.BuiltIn -> rest
153 V.Module _ -> rest
154 where
155 rest = concatMap ctors ps
158bound :: P.Canonical -> State (Set.Set String) ()
159bound pattern =
160 let
161 boundVars =
162 P.boundVarSet pattern
163 in
164 modify (\freeVars -> Set.difference freeVars boundVars)
167bindingReorder
168 :: (P.Canonical, Canonical.Expr)
169 -> State (Set.Set String) (P.Canonical, Canonical.Expr)
170bindingReorder (pattern, expr) =
171 do expr' <- reorder expr
172 bound pattern
173 mapM_ free (ctors pattern)
174 return (pattern, expr')
178-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS
181-- This also reorders the all of the sub-expressions in the Def list.
182buildDefGraph
183 :: [Canonical.Def]
184 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
185buildDefGraph defs =
186 do pdefsDeps <- mapM reorderAndGetDependencies defs
187 return $ realDeps (addKey pdefsDeps)
188 where
189 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
190 addKey =
191 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
193 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
194 variableToKey (Canonical.Definition _ pattern _ _, key, _) =
195 [ (var, key) | var <- P.boundVarList pattern ]
197 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
198 variableToKeyMap pdefsDeps =
199 Map.fromList (concatMap variableToKey pdefsDeps)
201 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
202 realDeps pdefsDeps =
203 map convert pdefsDeps
204 where
205 varDict =
206 variableToKeyMap pdefsDeps
208 convert (pdef, key, deps) =
209 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)
212reorderAndGetDependencies
213 :: Canonical.Def
214 -> State (Set.Set String) (Canonical.Def, [String])
215reorderAndGetDependencies (Canonical.Definition facts pattern expr mType) =
216 do globalFrees <- get
217 -- work in a fresh environment
218 put Set.empty
219 expr' <- reorder expr
220 localFrees <- get
221 -- merge with global frees
222 modify (Set.union globalFrees)
223 return (Canonical.Definition facts pattern expr' mType, Set.toList localFrees)