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