on 4940dfc help sign in
elk/compiler/src/Optimize/Expression.hs
475 lines 14 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.registerDebug name home region
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 <- destructCase 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 do obranches <- traverse (optimizeBranch temp) branches
144 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
146 Can.Accessor field ->
147 Names.registerField field (Opt.Accessor field)
149 Can.Access record field ->
150 do optRecord <- optimize cycle record
151 Names.registerField field (Opt.Access optRecord field)
153 Can.Update _ record updates ->
154 Names.registerFieldDict updates Opt.Update
155 <*> optimize cycle record
156 <*> traverse (optimizeUpdate cycle) updates
158 Can.Record fields ->
159 Names.registerFieldDict fields Opt.Record
160 <*> traverse (optimize cycle) fields
162 Can.Unit ->
163 Names.registerKernel N.utils Opt.Unit
165 Can.Tuple a b maybeC ->
166 Names.registerKernel N.utils Opt.Tuple
167 <*> optimize cycle a
168 <*> optimize cycle b
169 <*> traverse (optimize cycle) maybeC
171 Can.Shader _ src _ ->
172 pure (Opt.Shader src)
176-- UPDATE
179optimizeUpdate :: Cycle -> Can.FieldUpdate -> Names.Tracker Opt.Expr
180optimizeUpdate cycle (Can.FieldUpdate _ expr) =
181 optimize cycle expr
185-- DEFINITION
188optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
189optimizeDef cycle def body =
190 case def of
191 Can.Def (A.At _ name) args expr ->
192 optimizeDefHelp cycle name args expr body
194 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
195 optimizeDefHelp cycle name (map fst typedArgs) expr body
198optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
199optimizeDefHelp cycle name args expr body =
200 do oexpr <- optimize cycle expr
201 case args of
202 [] ->
203 pure $ Opt.Let (Opt.Def name oexpr) body
205 _ ->
206 do (argNames, destructors) <- destructArgs args
207 let ofunc = Opt.Function argNames (foldr Opt.Destruct oexpr destructors)
208 pure $ Opt.Let (Opt.Def name ofunc) body
212-- DESTRUCTURING
215destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [Opt.Destructor])
216destructArgs args =
217 do (argNames, destructorLists) <- unzip <$> traverse destruct args
218 return (argNames, concat destructorLists)
221destructCase :: N.Name -> Can.Pattern -> Names.Tracker [Opt.Destructor]
222destructCase rootName pattern =
223 reverse <$> destructHelp (Opt.Root rootName) pattern []
226destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Destructor])
227destruct pattern@(A.At _ ptrn) =
228 case ptrn of
229 Can.PVar name ->
230 pure (name, [])
232 Can.PAlias subPattern name ->
233 do revDs <- destructHelp (Opt.Root name) subPattern []
234 pure (name, reverse revDs)
236 _ ->
237 do name <- Names.generate
238 revDs <- destructHelp (Opt.Root name) pattern []
239 pure (name, reverse revDs)
242destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
243destructHelp path (A.At region pattern) revDs =
244 case pattern of
245 Can.PAnything ->
246 pure revDs
248 Can.PVar name ->
249 pure (Opt.Destructor name path : revDs)
251 Can.PRecord fields ->
252 let
253 toDestruct name =
254 Opt.Destructor name (Opt.Field name path)
255 in
256 Names.registerFieldList fields (map toDestruct fields ++ revDs)
258 Can.PAlias subPattern name ->
259 destructHelp (Opt.Root name) subPattern $
260 Opt.Destructor name path : revDs
262 Can.PUnit ->
263 pure revDs
265 Can.PTuple a b Nothing ->
266 destructTwo path a b revDs
268 Can.PTuple a b (Just c) ->
269 case path of
270 Opt.Root _ ->
271 destructHelp (Opt.Index Index.third path) c =<<
272 destructHelp (Opt.Index Index.second path) b =<<
273 destructHelp (Opt.Index Index.first path) a revDs
275 _ ->
276 do name <- Names.generate
277 let newRoot = Opt.Root name
278 destructHelp (Opt.Index Index.third newRoot) c =<<
279 destructHelp (Opt.Index Index.second newRoot) b =<<
280 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
282 Can.PList [] ->
283 pure revDs
285 Can.PList (hd:tl) ->
286 destructTwo path hd (A.At region (Can.PList tl)) revDs
288 Can.PCons hd tl ->
289 destructTwo path hd tl revDs
291 Can.PChr _ ->
292 pure revDs
294 Can.PStr _ ->
295 pure revDs
297 Can.PInt _ ->
298 pure revDs
300 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
301 case args of
302 [Can.PatternCtorArg _ _ arg] ->
303 case opts of
304 Can.Normal -> destructHelp (Opt.Index Index.first path) arg revDs
305 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
306 Can.Enum -> destructHelp (Opt.Index Index.first path) arg revDs
308 _ ->
309 case path of
310 Opt.Root _ ->
311 foldM (destructCtorArg path) revDs args
313 _ ->
314 do name <- Names.generate
315 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
318destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
319destructTwo path a b revDs =
320 case path of
321 Opt.Root _ ->
322 destructHelp (Opt.Index Index.second path) b =<<
323 destructHelp (Opt.Index Index.first path) a revDs
325 _ ->
326 do name <- Names.generate
327 let newRoot = Opt.Root name
328 destructHelp (Opt.Index Index.second newRoot) b =<<
329 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
332destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
333destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
334 destructHelp (Opt.Index index path) arg revDs
338-- TAIL CALL
341optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
342optimizePotentialTailCall cycle def =
343 case def of
344 Can.Def (A.At _ name) args expr ->
345 do (argNames, destructors) <- destructArgs args
346 toTailDef name argNames destructors <$>
347 optimizeTail cycle name argNames expr
349 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
350 do (argNames, destructors) <- destructArgs (map fst typedArgs)
351 toTailDef name argNames destructors <$>
352 optimizeTail cycle name argNames expr
355optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
356optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
357 case expression of
358 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
359 do oargs <- traverse (optimize cycle) args
360 if name == rootName
361 then
362 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
363 Index.LengthMatch pairs ->
364 pure $ Opt.TailCall name pairs
366 Index.LengthMismatch _ _ ->
367 do ofunc <- optimize cycle func
368 pure $ Opt.Call ofunc oargs
369 else
370 do ofunc <- optimize cycle func
371 pure $ Opt.Call ofunc oargs
373 Can.If branches finally ->
374 let
375 optimizeBranch (condition, branch) =
376 (,)
377 <$> optimize cycle condition
378 <*> optimizeTail cycle rootName argNames branch
379 in
380 Opt.If
381 <$> traverse optimizeBranch branches
382 <*> optimizeTail cycle rootName argNames finally
384 Can.Let def body ->
385 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
387 Can.LetRec defs body ->
388 case defs of
389 [def] ->
390 Opt.Let
391 <$> optimizePotentialTailCall cycle def
392 <*> optimizeTail cycle rootName argNames body
394 _ ->
395 do obody <- optimizeTail cycle rootName argNames body
396 foldM (\bod def -> optimizeDef cycle def bod) obody defs
398 Can.LetDestruct pattern expr body ->
399 do (dname, destructors) <- destruct pattern
400 oexpr <- optimize cycle expr
401 obody <- optimizeTail cycle dname argNames body
402 pure $
403 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
405 Can.Case expr branches ->
406 let
407 optimizeBranch root (Can.CaseBranch pattern branch) =
408 do destructors <- destructCase root pattern
409 obranch <- optimizeTail cycle rootName argNames branch
410 pure (pattern, foldr Opt.Destruct obranch destructors)
411 in
412 do temp <- Names.generate
413 oexpr <- optimize cycle expr
414 case oexpr of
415 Opt.VarLocal root ->
416 Case.optimize temp root <$> traverse (optimizeBranch root) branches
418 _ ->
419 do obranches <- traverse (optimizeBranch temp) branches
420 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
422 _ ->
423 optimize cycle locExpr
427-- DETECT TAIL CALLS
430toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
431toTailDef name argNames destructors body =
432 if hasTailCall body then
433 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
434 else
435 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
438hasTailCall :: Opt.Expr -> Bool
439hasTailCall expression =
440 case expression of
441 Opt.TailCall _ _ ->
442 True
444 Opt.If branches finally ->
445 hasTailCall finally || any (hasTailCall . snd) branches
447 Opt.Let _ body ->
448 hasTailCall body
450 Opt.Destruct _ body ->
451 hasTailCall body
453 Opt.Case _ _ decider jumps ->
454 decidecHasTailCall decider || any (hasTailCall . snd) jumps
456 _ ->
457 False
460decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
461decidecHasTailCall decider =
462 case decider of
463 Opt.Leaf choice ->
464 case choice of
465 Opt.Inline expr ->
466 hasTailCall expr
468 Opt.Jump _ ->
469 False
471 Opt.Chain _ success failure ->
472 decidecHasTailCall success || decidecHasTailCall failure
474 Opt.FanOut _ tests fallback ->
475 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests