on 3de162f help sign in
elk/compiler/src/Optimize/Expression.hs
481 lines 14 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Expression
4 ( optimize
5 , destructArgs
6 , optimizePotentialTailCall
7 )
8 where
11import Prelude hiding (cycle)
12import Control.Monad (foldM)
13import qualified Data.Set as Set
15import qualified AST.Canonical as Can
16import qualified AST.Optimized as Opt
17import qualified AST.Module.Name as ModuleName
18import qualified Data.Index as Index
19import qualified Elm.Name as N
20import qualified Optimize.Case as Case
21import qualified Optimize.Names as Names
22import qualified Reporting.Annotation as A
26-- OPTIMIZE
29type Cycle =
30 Set.Set N.Name
33optimize :: Cycle -> Can.Expr -> Names.Tracker Opt.Expr
34optimize cycle (A.At region expression) =
35 case expression of
36 Can.VarLocal name ->
37 pure (Opt.VarLocal name)
39 Can.VarTopLevel home name ->
40 if Set.member name cycle then
41 pure (Opt.VarCycle home name)
42 else
43 Names.registerGlobal home name
45 Can.VarKernel home name ->
46 Names.registerKernel home (Opt.VarKernel home name)
48 Can.VarForeign home name _ ->
49 Names.registerGlobal home name
51 Can.VarCtor opts home name index _ ->
52 Names.registerCtor home name index opts
54 Can.VarDebug home name _ ->
55 Names.registerDebug name home region
57 Can.VarOperator _ home name _ ->
58 Names.registerGlobal home name
60 Can.Chr chr ->
61 Names.registerKernel N.utils (Opt.Chr chr)
63 Can.Str str ->
64 pure (Opt.Str str)
66 Can.Int int ->
67 pure (Opt.Int int)
69 Can.Float float ->
70 pure (Opt.Float float)
72 Can.List entries ->
73 Names.registerKernel N.list Opt.List
74 <*> 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 <$> optimizePotentialTailCallDef 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 (A.At _ 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.PBool _ _ ->
301 pure revDs
303 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
304 case args of
305 [Can.PatternCtorArg _ _ arg] ->
306 case opts of
307 Can.Normal -> destructHelp (Opt.Index Index.first path) arg revDs
308 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
309 Can.Enum -> destructHelp (Opt.Index Index.first path) arg revDs
311 _ ->
312 case path of
313 Opt.Root _ ->
314 foldM (destructCtorArg path) revDs args
316 _ ->
317 do name <- Names.generate
318 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
321destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
322destructTwo path a b revDs =
323 case path of
324 Opt.Root _ ->
325 destructHelp (Opt.Index Index.second path) b =<<
326 destructHelp (Opt.Index Index.first path) a revDs
328 _ ->
329 do name <- Names.generate
330 let newRoot = Opt.Root name
331 destructHelp (Opt.Index Index.second newRoot) b =<<
332 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
335destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
336destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
337 destructHelp (Opt.Index index path) arg revDs
341-- TAIL CALL
344optimizePotentialTailCallDef :: Cycle -> Can.Def -> Names.Tracker Opt.Def
345optimizePotentialTailCallDef cycle def =
346 case def of
347 Can.Def (A.At _ name) args expr ->
348 optimizePotentialTailCall cycle name args expr
350 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
351 optimizePotentialTailCall cycle name (map fst typedArgs) expr
354optimizePotentialTailCall :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Names.Tracker Opt.Def
355optimizePotentialTailCall cycle name args expr =
356 do (argNames, destructors) <- destructArgs args
357 toTailDef name argNames destructors <$>
358 optimizeTail cycle name argNames expr
361optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
362optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
363 case expression of
364 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
365 do oargs <- traverse (optimize cycle) args
366 if name == rootName
367 then
368 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
369 Index.LengthMatch pairs ->
370 pure $ Opt.TailCall name pairs
372 Index.LengthMismatch _ _ ->
373 do ofunc <- optimize cycle func
374 pure $ Opt.Call ofunc oargs
375 else
376 do ofunc <- optimize cycle func
377 pure $ Opt.Call ofunc oargs
379 Can.If branches finally ->
380 let
381 optimizeBranch (condition, branch) =
382 (,)
383 <$> optimize cycle condition
384 <*> optimizeTail cycle rootName argNames branch
385 in
386 Opt.If
387 <$> traverse optimizeBranch branches
388 <*> optimizeTail cycle rootName argNames finally
390 Can.Let def body ->
391 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
393 Can.LetRec defs body ->
394 case defs of
395 [def] ->
396 Opt.Let
397 <$> optimizePotentialTailCallDef cycle def
398 <*> optimizeTail cycle rootName argNames body
400 _ ->
401 do obody <- optimizeTail cycle rootName argNames body
402 foldM (\bod def -> optimizeDef cycle def bod) obody defs
404 Can.LetDestruct pattern expr body ->
405 do (dname, destructors) <- destruct pattern
406 oexpr <- optimize cycle expr
407 obody <- optimizeTail cycle rootName argNames body
408 pure $
409 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
411 Can.Case expr branches ->
412 let
413 optimizeBranch root (Can.CaseBranch pattern branch) =
414 do destructors <- destructCase root pattern
415 obranch <- optimizeTail cycle rootName argNames branch
416 pure (pattern, foldr Opt.Destruct obranch destructors)
417 in
418 do temp <- Names.generate
419 oexpr <- optimize cycle expr
420 case oexpr of
421 Opt.VarLocal root ->
422 Case.optimize temp root <$> traverse (optimizeBranch root) branches
424 _ ->
425 do obranches <- traverse (optimizeBranch temp) branches
426 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
428 _ ->
429 optimize cycle locExpr
433-- DETECT TAIL CALLS
436toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
437toTailDef name argNames destructors body =
438 if hasTailCall body then
439 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
440 else
441 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
444hasTailCall :: Opt.Expr -> Bool
445hasTailCall expression =
446 case expression of
447 Opt.TailCall _ _ ->
448 True
450 Opt.If branches finally ->
451 hasTailCall finally || any (hasTailCall . snd) branches
453 Opt.Let _ body ->
454 hasTailCall body
456 Opt.Destruct _ body ->
457 hasTailCall body
459 Opt.Case _ _ decider jumps ->
460 decidecHasTailCall decider || any (hasTailCall . snd) jumps
462 _ ->
463 False
466decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
467decidecHasTailCall decider =
468 case decider of
469 Opt.Leaf choice ->
470 case choice of
471 Opt.Inline expr ->
472 hasTailCall expr
474 Opt.Jump _ ->
475 False
477 Opt.Chain _ success failure ->
478 decidecHasTailCall success || decidecHasTailCall failure
480 Opt.FanOut _ tests fallback ->
481 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests