git.iko.soy/elk
on f232da4 help sign in
elk/src/Optimize/Inline.hs
292 lines 6.7 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Inline (inline) where
5import Control.Arrow (second, (***))
6import Control.Monad (foldM)
7import qualified Data.List as List
8import qualified Data.Map as Map
9import Data.Text (Text)
11import AST.Expression.Optimized (Expr(..), Decider(..), Choice(..))
12import qualified AST.Expression.Optimized as Opt
13import qualified AST.Variable as Var
14import qualified Optimize.Environment as Env
18inline :: Map.Map Text Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
19inline rawSubs expression =
20 let
21 counts =
22 count expression
24 annotatedSubs =
25 Map.toList (Map.intersectionWith (,) counts rawSubs)
26 in
27 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
28 let inlinedExpr = replace subs expression
29 return $
30 if null defs then
31 inlinedExpr
33 else
34 Let (map (second Opt.Def) defs) inlinedExpr
37data Subs =
38 Subs
39 { _substitutions :: Map.Map Text Opt.Expr
40 , _definitions :: [(Text, Opt.Expr)]
41 }
44processSubs :: Subs -> (Text, (Int, Opt.Expr)) -> Env.Optimizer Subs
45processSubs (Subs subs defs) (name, (n, expr)) =
46 if n == 1 then
47 return $ Subs (Map.insert name expr subs) defs
49 else
50 do uniqueName <- Env.freshName
51 return $ Subs
52 (Map.insert name (Opt.Var (Var.local uniqueName)) subs)
53 ((uniqueName, expr) : defs)
57-- HELPERS
60deleteBatch :: [Text] -> Map.Map Text a -> Map.Map Text a
61deleteBatch names dict =
62 List.foldl' (flip Map.delete) dict names
66-- COUNT VARIABLE APPEARANCES
69count :: Opt.Expr -> Map.Map Text Int
70count expression =
71 let
72 count2 a b =
73 Map.unionWith (+) (count a) (count b)
75 countMany xs =
76 Map.unionsWith (+) (map count xs)
77 in
78 case expression of
79 Literal _ ->
80 Map.empty
82 Var (Var.Canonical home name) ->
83 case home of
84 Var.Local ->
85 Map.singleton name 1
87 Var.BuiltIn ->
88 Map.empty
90 Var.Module _ ->
91 Map.empty
93 Var.TopLevel _ ->
94 Map.empty
96 List exprs ->
97 countMany exprs
99 Binop _op left right ->
100 count2 left right
102 Function args body ->
103 deleteBatch args (count body)
105 Call func args ->
106 Map.unionWith (+) (count func) (countMany args)
108 TailCall _name _argNames args ->
109 countMany args
111 If branches finally ->
112 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
114 Let defs expr ->
115 let
116 countDef (name, def) =
117 case def of
118 Opt.Def body ->
119 Map.delete name (count body)
121 Opt.TailDef argNames body ->
122 deleteBatch (name:argNames) (count body)
124 exprCount =
125 deleteBatch (map fst defs) (count expr)
126 in
127 Map.unionsWith (+) (exprCount : map countDef defs)
129 Case _ decider jumps ->
130 let
131 countDecider dcdr =
132 case dcdr of
133 Leaf (Inline expr) ->
134 count expr
136 Leaf (Jump _) ->
137 Map.empty
139 Chain _ success failure ->
140 Map.unionWith (+) (countDecider success) (countDecider failure)
142 FanOut _path tests fallback ->
143 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
144 in
145 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
147 Ctor _tag values ->
148 countMany values
150 CtorAccess root _index ->
151 count root
153 Access record _field ->
154 count record
156 Update record fields ->
157 Map.unionWith (+) (count record) (countMany (map snd fields))
159 Record fields ->
160 countMany (map snd fields)
162 Cmd _ _ ->
163 Map.empty
165 Sub _ _ ->
166 Map.empty
168 OutgoingPort _ _ ->
169 Map.empty
171 IncomingPort _ _ ->
172 Map.empty
174 Program _ expr ->
175 count expr
177 GLShader _ _ _ ->
178 Map.empty
180 Crash _ _ maybeBranchProblem ->
181 maybe Map.empty count maybeBranchProblem
185-- REPLACE
188replace :: Map.Map Text Opt.Expr -> Opt.Expr -> Opt.Expr
189replace substitutions expression =
190 let
191 go = replace substitutions
192 in
193 case expression of
194 Literal _ ->
195 expression
197 Var (Var.Canonical Var.Local name) ->
198 maybe expression id (Map.lookup name substitutions)
200 Var _ ->
201 expression
203 List exprs ->
204 List (map go exprs)
206 Binop op left right ->
207 Binop op (go left) (go right)
209 Function args body ->
210 Function args (replace (deleteBatch args substitutions) body)
212 Call func args ->
213 Call (go func) (map go args)
215 TailCall name argNames args ->
216 TailCall name argNames (map go args)
218 If branches finally ->
219 If (map (go *** go) branches) (go finally)
221 Let defs expr ->
222 let
223 replaceDef (name, def) =
224 case def of
225 Opt.Def body ->
226 (,) name $ Opt.Def $
227 replace (Map.delete name substitutions) body
229 Opt.TailDef argNames body ->
230 (,) name $ Opt.TailDef argNames $
231 replace (deleteBatch (name:argNames) substitutions) body
233 boundNames =
234 map fst defs
235 in
236 Let
237 (map replaceDef defs)
238 (replace (deleteBatch boundNames substitutions) expr)
240 Case exprName decider jumps ->
241 let
242 goDecider dcdr =
243 case dcdr of
244 Leaf (Inline expr) ->
245 Leaf (Inline (go expr))
247 Leaf (Jump _target) ->
248 dcdr
250 Chain chain success failure ->
251 Chain chain (goDecider success) (goDecider failure)
253 FanOut path tests fallback ->
254 FanOut path (map (second goDecider) tests) (goDecider fallback)
255 in
256 Case exprName (goDecider decider) (map (second go) jumps)
258 Ctor tag values ->
259 Ctor tag (map go values)
261 CtorAccess root index ->
262 CtorAccess (go root) index
264 Access record field ->
265 Access (go record) field
267 Update record fields ->
268 Update (go record) (map (second go) fields)
270 Record fields ->
271 Record (map (second go) fields)
273 Cmd _ _ ->
274 expression
276 Sub _ _ ->
277 expression
279 OutgoingPort _ _ ->
280 expression
282 IncomingPort _ _ ->
283 expression
285 Program kind expr ->
286 Program kind (go expr)
288 GLShader _ _ _ ->
289 expression
291 Crash home region maybeBranchProblem ->
292 Crash home region (fmap go maybeBranchProblem)