git.iko.soy/elk
on efb86bc help sign in
elk/src/Optimize/Inline.hs
294 lines 7.1 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.General (PortImpl(..))
10import AST.Expression.Optimized (Expr(..), Decider(..), Choice(..))
11import qualified AST.Expression.Optimized as Opt
12import qualified AST.Variable as Var
13import qualified Optimize.Environment as Env
17inline :: Map.Map String Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
18inline rawSubs expression =
19 let
20 counts =
21 count expression
23 annotatedSubs =
24 Map.toList (Map.intersectionWith (,) counts rawSubs)
25 in
26 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
27 let inlinedExpr = replace subs expression
28 return $
29 if null defs then
30 inlinedExpr
32 else
33 Let (map (\(name, expr) -> Opt.Def Opt.dummyFacts name expr) defs) inlinedExpr
36data Subs = Subs
37 { _substitutions :: Map.Map String Opt.Expr
38 , _definitions :: [(String, Opt.Expr)]
39 }
42processSubs :: Subs -> (String, (Int, Opt.Expr)) -> Env.Optimizer Subs
43processSubs (Subs subs defs) (name, (n, expr)) =
44 if n == 1 then
45 return $ Subs (Map.insert name expr subs) defs
47 else
48 do uniqueName <- Env.freshName
49 return $ Subs
50 (Map.insert name (Opt.Var (Var.local uniqueName)) subs)
51 ((uniqueName, expr) : defs)
53-- HELPERS
55deleteBatch :: [String] -> Map.Map String a -> Map.Map String a
56deleteBatch names dict =
57 List.foldl' (flip Map.delete) dict names
60getDefName :: Opt.Def -> String
61getDefName def =
62 case def of
63 Opt.Def _ name _ ->
64 name
66 Opt.TailDef _ name _ _ ->
67 name
70-- COUNT VARIABLE APPEARANCES
72count :: Opt.Expr -> Map.Map String Int
73count expression =
74 let
75 count2 a b =
76 Map.unionWith (+) (count a) (count b)
78 countMany xs =
79 Map.unionsWith (+) (map count xs)
80 in
81 case expression of
82 Literal _ ->
83 Map.empty
85 Var (Var.Canonical home name) ->
86 case home of
87 Var.Local ->
88 Map.singleton name 1
90 Var.BuiltIn ->
91 Map.empty
93 Var.Module _ ->
94 Map.empty
96 Var.TopLevel _ ->
97 Map.empty
99 Range lo hi ->
100 count2 lo hi
102 ExplicitList exprs ->
103 countMany exprs
105 Binop _op left right ->
106 count2 left right
108 Function args body ->
109 deleteBatch args (count body)
111 Call func args ->
112 Map.unionWith (+) (count func) (countMany args)
114 TailCall _name _argNames args ->
115 countMany args
117 If branches finally ->
118 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
120 Let defs expr ->
121 let
122 countDef def =
123 case def of
124 Opt.Def _ name body ->
125 Map.delete name (count body)
127 Opt.TailDef _ name argNames body ->
128 deleteBatch (name:argNames) (count body)
130 exprCount =
131 deleteBatch (map getDefName defs) (count expr)
132 in
133 Map.unionsWith (+) (exprCount : map countDef defs)
135 Case _ decider jumps ->
136 let
137 countDecider dcdr =
138 case dcdr of
139 Leaf (Inline expr) ->
140 count expr
142 Leaf (Jump _) ->
143 Map.empty
145 Chain _ success failure ->
146 Map.unionWith (+) (countDecider success) (countDecider failure)
148 FanOut _path tests fallback ->
149 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
150 in
151 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
153 Data _tag values ->
154 countMany values
156 DataAccess root _index ->
157 count root
159 Access record _field ->
160 count record
162 Update record fields ->
163 Map.unionWith (+) (count record) (countMany (map snd fields))
165 Record fields ->
166 countMany (map snd fields)
168 Port impl ->
169 case impl of
170 In _ _ ->
171 Map.empty
173 Out _name expr _tipe ->
174 count expr
176 Task _name expr _tipe ->
177 count expr
179 GLShader _ _ _ ->
180 Map.empty
182 Crash _ _ maybeBranchProblem ->
183 maybe Map.empty count maybeBranchProblem
187-- REPLACE
190replace :: Map.Map String Opt.Expr -> Opt.Expr -> Opt.Expr
191replace substitutions expression =
192 let
193 go = replace substitutions
194 in
195 case expression of
196 Literal _ ->
197 expression
199 Var (Var.Canonical Var.Local name) ->
200 maybe expression id (Map.lookup name substitutions)
202 Var _ ->
203 expression
205 Range lo hi ->
206 Range (go lo) (go hi)
208 ExplicitList exprs ->
209 ExplicitList (map go exprs)
211 Binop op left right ->
212 Binop op (go left) (go right)
214 Function args body ->
215 Function args (replace (deleteBatch args substitutions) body)
217 Call func args ->
218 Call (go func) (map go args)
220 TailCall name argNames args ->
221 TailCall name argNames (map go args)
223 If branches finally ->
224 If (map (go *** go) branches) (go finally)
226 Let defs expr ->
227 let
228 replaceDef def =
229 case def of
230 Opt.Def facts name body ->
231 Opt.Def facts name
232 (replace (Map.delete name substitutions) body)
234 Opt.TailDef facts name argNames body ->
235 Opt.TailDef facts name argNames
236 (replace (deleteBatch (name:argNames) substitutions) body)
238 boundNames =
239 map getDefName defs
240 in
241 Let
242 (map replaceDef defs)
243 (replace (deleteBatch boundNames substitutions) expr)
245 Case exprName decider jumps ->
246 let
247 goDecider dcdr =
248 case dcdr of
249 Leaf (Inline expr) ->
250 Leaf (Inline (go expr))
252 Leaf (Jump _target) ->
253 dcdr
255 Chain chain success failure ->
256 Chain chain (goDecider success) (goDecider failure)
258 FanOut path tests fallback ->
259 FanOut path (map (second goDecider) tests) (goDecider fallback)
260 in
261 Case exprName (goDecider decider) (map (second go) jumps)
263 Data tag values ->
264 Data tag (map go values)
266 DataAccess root index ->
267 DataAccess (go root) index
269 Access record field ->
270 Access (go record) field
272 Update record fields ->
273 Update (go record) (map (second go) fields)
275 Record fields ->
276 Record (map (second go) fields)
278 Port impl ->
279 Port $
280 case impl of
281 In _ _ ->
282 impl
284 Out name expr tipe ->
285 Out name (go expr) tipe
287 Task name expr tipe ->
288 Task name (go expr) tipe
290 GLShader _ _ _ ->
291 expression
293 Crash home region maybeBranchProblem ->
294 Crash home region (fmap go maybeBranchProblem)