git.iko.soy/elk
on 54dfb21 help sign in
elk/src/Canonicalize/Sort.hs
219 lines 6.3 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'(..), PortImpl(..))
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
17-- STARTING POINT
19definitions :: Canonical.Expr -> Canonical.Expr
20definitions expression =
21 evalState (reorder expression) Set.empty
24free :: String -> State (Set.Set String) ()
25free x =
26 modify (Set.insert x)
29freeIfLocal :: V.Canonical -> State (Set.Set String) ()
30freeIfLocal (V.Canonical home name) =
31 case home of
32 V.Local -> free name
33 V.TopLevel _ -> free name
34 V.BuiltIn -> return ()
35 V.Module _ -> return ()
38-- REORDER EXPRESSIONS
40reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
41reorder (A.A ann expression) =
42 A.A ann <$>
43 case expression of
44 -- Be careful adding and restricting freeVars
45 Var var ->
46 do freeIfLocal var
47 return expression
49 Lambda pattern body ->
50 uncurry Lambda <$> bindingReorder (pattern,body)
52 Binop op leftExpr rightExpr ->
53 do freeIfLocal op
54 Binop op <$> reorder leftExpr <*> reorder rightExpr
56 Case expr cases ->
57 Case <$> reorder expr <*> mapM bindingReorder cases
59 Data name exprs ->
60 do free name
61 Data name <$> mapM reorder exprs
63 -- Just pipe the reorder though
64 Literal _ ->
65 return expression
67 Range lowExpr highExpr ->
68 Range <$> reorder lowExpr <*> reorder highExpr
70 ExplicitList es ->
71 ExplicitList <$> mapM reorder es
73 App func arg ->
74 App <$> reorder func <*> reorder arg
76 If branches finally ->
77 If
78 <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
79 <*> reorder finally
81 Access record field ->
82 Access
83 <$> reorder record
84 <*> return field
86 Update record fields ->
87 Update
88 <$> reorder record
89 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
91 Record fields ->
92 Record
93 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
95 GLShader _ _ _ ->
96 return expression
98 Port impl ->
99 Port <$>
100 case impl of
101 In _ _ ->
102 return impl
104 Out name expr tipe ->
105 (\e -> Out name e tipe) <$> reorder expr
107 Task name expr tipe ->
108 (\e -> Task name e tipe) <$> reorder expr
110 -- Actually do some reordering
111 Let defs body ->
112 do body' <- reorder body
114 -- Sort defs into strongly connected components.This
115 -- allows the programmer to write definitions in whatever
116 -- order they please, we can still define things in order
117 -- and generalize polymorphic functions when appropriate.
118 sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
119 let defss = map Graph.flattenSCC sccs
121 -- remove let-bound variables from the context
122 forM_ defs $ \(Canonical.Definition _ pattern _ _) -> do
123 bound pattern
124 mapM free (ctors pattern)
126 let (A.A _ let') =
127 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
129 return let'
132ctors :: P.CanonicalPattern -> [String]
133ctors (A.A _ pattern) =
134 case pattern of
135 P.Var _ ->
136 []
138 P.Alias _ p ->
139 ctors p
141 P.Record _ ->
142 []
144 P.Anything ->
145 []
147 P.Literal _ ->
148 []
150 P.Data (V.Canonical home name) ps ->
151 case home of
152 V.Local -> name : rest
153 V.TopLevel _ -> name : rest
154 V.BuiltIn -> rest
155 V.Module _ -> rest
156 where
157 rest = concatMap ctors ps
160bound :: P.CanonicalPattern -> State (Set.Set String) ()
161bound pattern =
162 let boundVars = P.boundVarSet pattern
163 in
164 modify (\freeVars -> Set.difference freeVars boundVars)
167bindingReorder
168 :: (P.CanonicalPattern, Canonical.Expr)
169 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
170bindingReorder (pattern, expr) =
171 do expr' <- reorder expr
172 bound pattern
173 mapM_ free (ctors pattern)
174 return (pattern, expr')
177-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS
179-- This also reorders the all of the sub-expressions in the Def list.
180buildDefGraph
181 :: [Canonical.Def]
182 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
183buildDefGraph defs =
184 do pdefsDeps <- mapM reorderAndGetDependencies defs
185 return $ realDeps (addKey pdefsDeps)
186 where
187 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
188 addKey =
189 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
191 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
192 variableToKey (Canonical.Definition _ pattern _ _, key, _) =
193 [ (var, key) | var <- P.boundVarList pattern ]
195 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
196 variableToKeyMap pdefsDeps =
197 Map.fromList (concatMap variableToKey pdefsDeps)
199 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
200 realDeps pdefsDeps =
201 map convert pdefsDeps
202 where
203 varDict = variableToKeyMap pdefsDeps
204 convert (pdef, key, deps) =
205 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)
208reorderAndGetDependencies
209 :: Canonical.Def
210 -> State (Set.Set String) (Canonical.Def, [String])
211reorderAndGetDependencies (Canonical.Definition facts pattern expr mType) =
212 do globalFrees <- get
213 -- work in a fresh environment
214 put Set.empty
215 expr' <- reorder expr
216 localFrees <- get
217 -- merge with global frees
218 modify (Set.union globalFrees)
219 return (Canonical.Definition facts pattern expr' mType, Set.toList localFrees)