git.iko.soy/elk
on fa8ade2 help sign in
elk/src/Optimize/Inline.hs
305 lines 7.0 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Inline (inline) where
4import Control.Arrow (second, (***))
5import Control.Monad (foldM)
6import qualified Data.List as List
7import qualified Data.Map as Map
9import AST.Expression.Optimized (Expr(..), Decider(..), Choice(..))
10import qualified AST.Expression.Optimized as Opt
11import qualified AST.Variable as Var
12import qualified Optimize.Environment as Env
16inline :: Map.Map String Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
17inline rawSubs expression =
18 let
19 counts =
20 count expression
22 annotatedSubs =
23 Map.toList (Map.intersectionWith (,) counts rawSubs)
24 in
25 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
26 let inlinedExpr = replace subs expression
27 return $
28 if null defs then
29 inlinedExpr
31 else
32 Let (map (\(name, expr) -> Opt.Def Opt.dummyFacts name expr) defs) inlinedExpr
35data Subs = Subs
36 { _substitutions :: Map.Map String Opt.Expr
37 , _definitions :: [(String, Opt.Expr)]
38 }
41processSubs :: Subs -> (String, (Int, Opt.Expr)) -> Env.Optimizer Subs
42processSubs (Subs subs defs) (name, (n, expr)) =
43 if n == 1 then
44 return $ Subs (Map.insert name expr subs) defs
46 else
47 do uniqueName <- Env.freshName
48 return $ Subs
49 (Map.insert name (Opt.Var (Var.local uniqueName)) subs)
50 ((uniqueName, expr) : defs)
54-- HELPERS
57deleteBatch :: [String] -> Map.Map String a -> Map.Map String a
58deleteBatch names dict =
59 List.foldl' (flip Map.delete) dict names
62getDefName :: Opt.Def -> String
63getDefName def =
64 case def of
65 Opt.Def _ name _ ->
66 name
68 Opt.TailDef _ name _ _ ->
69 name
73-- COUNT VARIABLE APPEARANCES
76count :: Opt.Expr -> Map.Map String Int
77count expression =
78 let
79 count2 a b =
80 Map.unionWith (+) (count a) (count b)
82 countMany xs =
83 Map.unionsWith (+) (map count xs)
84 in
85 case expression of
86 Literal _ ->
87 Map.empty
89 Var (Var.Canonical home name) ->
90 case home of
91 Var.Local ->
92 Map.singleton name 1
94 Var.BuiltIn ->
95 Map.empty
97 Var.Module _ ->
98 Map.empty
100 Var.TopLevel _ ->
101 Map.empty
103 Range lo hi ->
104 count2 lo hi
106 ExplicitList exprs ->
107 countMany exprs
109 Binop _op left right ->
110 count2 left right
112 Function args body ->
113 deleteBatch args (count body)
115 Call func args ->
116 Map.unionWith (+) (count func) (countMany args)
118 TailCall _name _argNames args ->
119 countMany args
121 If branches finally ->
122 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
124 Let defs expr ->
125 let
126 countDef def =
127 case def of
128 Opt.Def _ name body ->
129 Map.delete name (count body)
131 Opt.TailDef _ name argNames body ->
132 deleteBatch (name:argNames) (count body)
134 exprCount =
135 deleteBatch (map getDefName defs) (count expr)
136 in
137 Map.unionsWith (+) (exprCount : map countDef defs)
139 Case _ decider jumps ->
140 let
141 countDecider dcdr =
142 case dcdr of
143 Leaf (Inline expr) ->
144 count expr
146 Leaf (Jump _) ->
147 Map.empty
149 Chain _ success failure ->
150 Map.unionWith (+) (countDecider success) (countDecider failure)
152 FanOut _path tests fallback ->
153 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
154 in
155 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
157 Data _tag values ->
158 countMany values
160 DataAccess root _index ->
161 count root
163 Access record _field ->
164 count record
166 Update record fields ->
167 Map.unionWith (+) (count record) (countMany (map snd fields))
169 Record fields ->
170 countMany (map snd fields)
172 Cmd _ ->
173 Map.empty
175 Sub _ ->
176 Map.empty
178 OutgoingPort _ _ ->
179 Map.empty
181 IncomingPort _ _ ->
182 Map.empty
184 Program _ expr ->
185 count expr
187 GLShader _ _ _ ->
188 Map.empty
190 Crash _ _ maybeBranchProblem ->
191 maybe Map.empty count maybeBranchProblem
195-- REPLACE
198replace :: Map.Map String Opt.Expr -> Opt.Expr -> Opt.Expr
199replace substitutions expression =
200 let
201 go = replace substitutions
202 in
203 case expression of
204 Literal _ ->
205 expression
207 Var (Var.Canonical Var.Local name) ->
208 maybe expression id (Map.lookup name substitutions)
210 Var _ ->
211 expression
213 Range lo hi ->
214 Range (go lo) (go hi)
216 ExplicitList exprs ->
217 ExplicitList (map go exprs)
219 Binop op left right ->
220 Binop op (go left) (go right)
222 Function args body ->
223 Function args (replace (deleteBatch args substitutions) body)
225 Call func args ->
226 Call (go func) (map go args)
228 TailCall name argNames args ->
229 TailCall name argNames (map go args)
231 If branches finally ->
232 If (map (go *** go) branches) (go finally)
234 Let defs expr ->
235 let
236 replaceDef def =
237 case def of
238 Opt.Def facts name body ->
239 Opt.Def facts name
240 (replace (Map.delete name substitutions) body)
242 Opt.TailDef facts name argNames body ->
243 Opt.TailDef facts name argNames
244 (replace (deleteBatch (name:argNames) substitutions) body)
246 boundNames =
247 map getDefName defs
248 in
249 Let
250 (map replaceDef defs)
251 (replace (deleteBatch boundNames substitutions) expr)
253 Case exprName decider jumps ->
254 let
255 goDecider dcdr =
256 case dcdr of
257 Leaf (Inline expr) ->
258 Leaf (Inline (go expr))
260 Leaf (Jump _target) ->
261 dcdr
263 Chain chain success failure ->
264 Chain chain (goDecider success) (goDecider failure)
266 FanOut path tests fallback ->
267 FanOut path (map (second goDecider) tests) (goDecider fallback)
268 in
269 Case exprName (goDecider decider) (map (second go) jumps)
271 Data tag values ->
272 Data tag (map go values)
274 DataAccess root index ->
275 DataAccess (go root) index
277 Access record field ->
278 Access (go record) field
280 Update record fields ->
281 Update (go record) (map (second go) fields)
283 Record fields ->
284 Record (map (second go) fields)
286 Cmd _ ->
287 expression
289 Sub _ ->
290 expression
292 OutgoingPort _ _ ->
293 expression
295 IncomingPort _ _ ->
296 expression
298 Program kind expr ->
299 Program kind (go expr)
301 GLShader _ _ _ ->
302 expression
304 Crash home region maybeBranchProblem ->
305 Crash home region (fmap go maybeBranchProblem)