git.iko.soy/elk
on 7de84bc help sign in
elk/src/Optimize/Inline.hs
287 lines 6.8 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 GLShader _ _ _ ->
179 Map.empty
181 Crash _ _ maybeBranchProblem ->
182 maybe Map.empty count maybeBranchProblem
186-- REPLACE
189replace :: Map.Map String Opt.Expr -> Opt.Expr -> Opt.Expr
190replace substitutions expression =
191 let
192 go = replace substitutions
193 in
194 case expression of
195 Literal _ ->
196 expression
198 Var (Var.Canonical Var.Local name) ->
199 maybe expression id (Map.lookup name substitutions)
201 Var _ ->
202 expression
204 Range lo hi ->
205 Range (go lo) (go hi)
207 ExplicitList exprs ->
208 ExplicitList (map go exprs)
210 Binop op left right ->
211 Binop op (go left) (go right)
213 Function args body ->
214 Function args (replace (deleteBatch args substitutions) body)
216 Call func args ->
217 Call (go func) (map go args)
219 TailCall name argNames args ->
220 TailCall name argNames (map go args)
222 If branches finally ->
223 If (map (go *** go) branches) (go finally)
225 Let defs expr ->
226 let
227 replaceDef def =
228 case def of
229 Opt.Def facts name body ->
230 Opt.Def facts name
231 (replace (Map.delete name substitutions) body)
233 Opt.TailDef facts name argNames body ->
234 Opt.TailDef facts name argNames
235 (replace (deleteBatch (name:argNames) substitutions) body)
237 boundNames =
238 map getDefName defs
239 in
240 Let
241 (map replaceDef defs)
242 (replace (deleteBatch boundNames substitutions) expr)
244 Case exprName decider jumps ->
245 let
246 goDecider dcdr =
247 case dcdr of
248 Leaf (Inline expr) ->
249 Leaf (Inline (go expr))
251 Leaf (Jump _target) ->
252 dcdr
254 Chain chain success failure ->
255 Chain chain (goDecider success) (goDecider failure)
257 FanOut path tests fallback ->
258 FanOut path (map (second goDecider) tests) (goDecider fallback)
259 in
260 Case exprName (goDecider decider) (map (second go) jumps)
262 Data tag values ->
263 Data tag (map go values)
265 DataAccess root index ->
266 DataAccess (go root) index
268 Access record field ->
269 Access (go record) field
271 Update record fields ->
272 Update (go record) (map (second go) fields)
274 Record fields ->
275 Record (map (second go) fields)
277 Cmd _ ->
278 expression
280 Sub _ ->
281 expression
283 GLShader _ _ _ ->
284 expression
286 Crash home region maybeBranchProblem ->
287 Crash home region (fmap go maybeBranchProblem)