git.iko.soy/elk
on 3492239 help sign in
elk/src/Optimize/Expression.hs
451 lines 13 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Expression
4 ( Cycle
5 , optimize
6 , destructArgs
7 , optimizePotentialTailCall
8 )
9 where
12import Prelude hiding (cycle)
13import Control.Monad (foldM)
14import qualified Data.Set as Set
16import qualified AST.Canonical as Can
17import qualified AST.Optimized as Opt
18import qualified AST.Module.Name as ModuleName
19import qualified Data.Index as Index
20import qualified Elm.Name as N
21import qualified Optimize.Case as Case
22import qualified Optimize.Names as Names
23import qualified Reporting.Annotation as A
27-- OPTIMIZE
30type Cycle =
31 Set.Set N.Name
34optimize :: Cycle -> Can.Expr -> Names.Tracker Opt.Expr
35optimize cycle (A.At region expression) =
36 case expression of
37 Can.VarLocal name ->
38 pure (Opt.VarLocal name)
40 Can.VarTopLevel home name ->
41 if Set.member name cycle then
42 pure (Opt.VarCycle home name)
43 else
44 Names.registerGlobal home name
46 Can.VarKernel home name ->
47 Names.registerKernel home (Opt.VarKernel home name)
49 Can.VarForeign home name _ ->
50 Names.registerGlobal home name
52 Can.VarCtor opts home name index _ ->
53 Names.registerCtor home name index opts
55 Can.VarDebug home name _ ->
56 Names.registerKernel N.debug (Opt.VarDebug name home region Nothing)
58 Can.VarOperator _ home name _ ->
59 Names.registerGlobal home name
61 Can.Chr chr ->
62 Names.registerKernel N.utils (Opt.Chr chr)
64 Can.Str str ->
65 pure (Opt.Str str)
67 Can.Int int ->
68 pure (Opt.Int int)
70 Can.Float float ->
71 pure (Opt.Float float)
73 Can.List entries ->
74 Opt.List <$> traverse (optimize cycle) entries
76 Can.Negate expr ->
77 do func <- Names.registerGlobal ModuleName.basics N.negate
78 arg <- optimize cycle expr
79 pure $ Opt.Call func [arg]
81 Can.Binop _ home name _ left right ->
82 do optFunc <- Names.registerGlobal home name
83 optLeft <- optimize cycle left
84 optRight <- optimize cycle right
85 return (Opt.Call optFunc [optLeft, optRight])
87 Can.Lambda args body ->
88 do (argNames, destructors) <- destructArgs args
89 obody <- optimize cycle body
90 pure $ Opt.Function argNames (foldr Opt.Destruct obody destructors)
92 Can.Call func args ->
93 Opt.Call
94 <$> optimize cycle func
95 <*> traverse (optimize cycle) args
97 Can.If branches finally ->
98 let
99 optimizeBranch (condition, branch) =
100 (,)
101 <$> optimize cycle condition
102 <*> optimize cycle branch
103 in
104 Opt.If
105 <$> traverse optimizeBranch branches
106 <*> optimize cycle finally
108 Can.Let def body ->
109 optimizeDef cycle def =<< optimize cycle body
111 Can.LetRec defs body ->
112 case defs of
113 [def] ->
114 Opt.Let
115 <$> optimizePotentialTailCall cycle def
116 <*> optimize cycle body
118 _ ->
119 do obody <- optimize cycle body
120 foldM (\bod def -> optimizeDef cycle def bod) obody defs
122 Can.LetDestruct pattern expr body ->
123 do (name, destructs) <- destruct pattern
124 oexpr <- optimize cycle expr
125 obody <- optimize cycle body
126 pure $
127 Opt.Let (Opt.Def name oexpr) (foldr Opt.Destruct obody destructs)
129 Can.Case expr branches ->
130 let
131 optimizeBranch root (Can.CaseBranch pattern branch) =
132 do destructors <- destructHelp (Opt.Root root) pattern []
133 obranch <- optimize cycle branch
134 pure (pattern, foldr Opt.Destruct obranch destructors)
135 in
136 do temp <- Names.generate
137 oexpr <- optimize cycle expr
138 case oexpr of
139 Opt.VarLocal root ->
140 Case.optimize temp root <$> traverse (optimizeBranch root) branches
142 _ ->
143 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
145 Can.Accessor field ->
146 Names.registerField field (Opt.Accessor field)
148 Can.Access record field ->
149 do optRecord <- optimize cycle record
150 Names.registerField field (Opt.Access optRecord field)
152 Can.Update record updates ->
153 Names.registerFieldDict updates Opt.Update
154 <*> optimize cycle record
155 <*> traverse (optimize cycle) updates
157 Can.Record fields ->
158 Names.registerFieldDict fields Opt.Record
159 <*> traverse (optimize cycle) fields
161 Can.Unit ->
162 Names.registerKernel N.utils Opt.Unit
164 Can.Tuple a b maybeC ->
165 Names.registerKernel N.utils Opt.Tuple
166 <*> optimize cycle a
167 <*> optimize cycle b
168 <*> traverse (optimize cycle) maybeC
170 Can.Shader _ src _ ->
171 pure (Opt.Shader src)
175-- DEFINITION
178optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
179optimizeDef cycle def body =
180 case def of
181 Can.Def (A.At _ name) args expr ->
182 optimizeDefHelp cycle name args expr body
184 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
185 optimizeDefHelp cycle name (map fst typedArgs) expr body
188optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
189optimizeDefHelp cycle name args expr body =
190 do oexpr <- optimize cycle expr
191 case args of
192 [] ->
193 pure $ Opt.Let (Opt.Def name oexpr) body
195 _ ->
196 do (argNames, destructors) <- destructArgs args
197 let ofunc = Opt.Function argNames (foldr Opt.Destruct oexpr destructors)
198 pure $ Opt.Let (Opt.Def name ofunc) body
202-- DESTRUCTURING
205destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [Opt.Destructor])
206destructArgs args =
207 do (argNames, destructorLists) <- unzip <$> traverse destruct args
208 return (argNames, concat destructorLists)
211destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Destructor])
212destruct pattern@(A.At _ ptrn) =
213 case ptrn of
214 Can.PVar name ->
215 pure (name, [])
217 Can.PAlias subPattern name ->
218 do revDs <- destructHelp (Opt.Root name) subPattern []
219 pure (name, reverse revDs)
221 _ ->
222 do name <- Names.generate
223 revDs <- destructHelp (Opt.Root name) pattern []
224 pure (name, reverse revDs)
227destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
228destructHelp path (A.At region pattern) revDs =
229 case pattern of
230 Can.PAnything ->
231 pure revDs
233 Can.PVar name ->
234 pure (Opt.Destructor name path : revDs)
236 Can.PRecord fields ->
237 let
238 toDestruct name =
239 Opt.Destructor name (Opt.Field name path)
240 in
241 Names.registerFieldList fields (map toDestruct fields ++ revDs)
243 Can.PAlias subPattern name ->
244 destructHelp (Opt.Root name) subPattern $
245 Opt.Destructor name path : revDs
247 Can.PUnit ->
248 pure revDs
250 Can.PTuple a b Nothing ->
251 destructTwo path a b revDs
253 Can.PTuple a b (Just c) ->
254 case path of
255 Opt.Root _ ->
256 destructHelp (Opt.Index Index.first path) a =<<
257 destructHelp (Opt.Index Index.second path) b =<<
258 destructHelp (Opt.Index Index.third path) c revDs
260 _ ->
261 do name <- Names.generate
262 let newRoot = Opt.Root name
263 destructHelp (Opt.Index Index.first newRoot) a =<<
264 destructHelp (Opt.Index Index.second newRoot) b =<<
265 destructHelp (Opt.Index Index.third newRoot) c (Opt.Destructor name path : revDs)
267 Can.PList [] ->
268 pure revDs
270 Can.PList (hd:tl) ->
271 destructTwo path hd (A.At region (Can.PList tl)) revDs
273 Can.PCons hd tl ->
274 destructTwo path hd tl revDs
276 Can.PChr _ ->
277 pure revDs
279 Can.PStr _ ->
280 pure revDs
282 Can.PInt _ ->
283 pure revDs
285 Can.PCtor _ _ _ _ _ args ->
286 case path of
287 Opt.Root _ ->
288 foldM (destructCtorArg path) revDs args
290 _ ->
291 do name <- Names.generate
292 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
295destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
296destructTwo path a b revDs =
297 case path of
298 Opt.Root _ ->
299 destructHelp (Opt.Index Index.first path) a =<<
300 destructHelp (Opt.Index Index.second path) b revDs
302 _ ->
303 do name <- Names.generate
304 let newRoot = Opt.Root name
305 destructHelp (Opt.Index Index.first newRoot) a =<<
306 destructHelp (Opt.Index Index.second newRoot) b (Opt.Destructor name path : revDs)
309destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
310destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
311 destructHelp (Opt.Index index path) arg revDs
315-- TAIL CALL
318optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
319optimizePotentialTailCall cycle def =
320 case def of
321 Can.Def (A.At _ name) args expr ->
322 do (argNames, destructors) <- destructArgs args
323 toTailDef name argNames destructors <$>
324 optimizeTail cycle name argNames expr
326 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
327 do (argNames, destructors) <- destructArgs (map fst typedArgs)
328 toTailDef name argNames destructors <$>
329 optimizeTail cycle name argNames expr
332optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
333optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
334 case expression of
335 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
336 do oargs <- traverse (optimize cycle) args
337 if name == rootName
338 then
339 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
340 Index.LengthMatch pairs ->
341 pure $ Opt.TailCall name pairs
343 Index.LengthMismatch _ _ ->
344 do ofunc <- optimize cycle func
345 pure $ Opt.Call ofunc oargs
346 else
347 do ofunc <- optimize cycle func
348 pure $ Opt.Call ofunc oargs
350 Can.If branches finally ->
351 let
352 optimizeBranch (condition, branch) =
353 (,)
354 <$> optimize cycle condition
355 <*> optimizeTail cycle rootName argNames branch
356 in
357 Opt.If
358 <$> traverse optimizeBranch branches
359 <*> optimizeTail cycle rootName argNames finally
361 Can.Let def body ->
362 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
364 Can.LetRec defs body ->
365 case defs of
366 [def] ->
367 Opt.Let
368 <$> optimizePotentialTailCall cycle def
369 <*> optimizeTail cycle rootName argNames body
371 _ ->
372 do obody <- optimizeTail cycle rootName argNames body
373 foldM (\bod def -> optimizeDef cycle def bod) obody defs
375 Can.LetDestruct pattern expr body ->
376 do (dname, destructors) <- destruct pattern
377 oexpr <- optimize cycle expr
378 obody <- optimizeTail cycle dname argNames body
379 pure $
380 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
382 Can.Case expr branches ->
383 let
384 optimizeBranch root (Can.CaseBranch pattern branch) =
385 do destructors <- destructHelp (Opt.Root root) pattern []
386 obranch <- optimizeTail cycle rootName argNames branch
387 pure (pattern, foldr Opt.Destruct obranch destructors)
388 in
389 do temp <- Names.generate
390 oexpr <- optimize cycle expr
391 case oexpr of
392 Opt.VarLocal root ->
393 Case.optimize temp root <$> traverse (optimizeBranch root) branches
395 _ ->
396 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
398 _ ->
399 optimize cycle locExpr
403-- DETECT TAIL CALLS
406toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
407toTailDef name argNames destructors body =
408 if hasTailCall body then
409 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
410 else
411 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
414hasTailCall :: Opt.Expr -> Bool
415hasTailCall expression =
416 case expression of
417 Opt.TailCall _ _ ->
418 True
420 Opt.If branches finally ->
421 hasTailCall finally || any (hasTailCall . snd) branches
423 Opt.Let _ body ->
424 hasTailCall body
426 Opt.Destruct _ body ->
427 hasTailCall body
429 Opt.Case _ _ decider jumps ->
430 decidecHasTailCall decider || any (hasTailCall . snd) jumps
432 _ ->
433 False
436decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
437decidecHasTailCall decider =
438 case decider of
439 Opt.Leaf choice ->
440 case choice of
441 Opt.Inline expr ->
442 hasTailCall expr
444 Opt.Jump _ ->
445 False
447 Opt.Chain _ success failure ->
448 decidecHasTailCall success || decidecHasTailCall failure
450 Opt.FanOut _ tests fallback ->
451 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests