on 233cfc8 help sign in
elk/compiler/src/Optimize/Expression.hs
476 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 Names.registerKernel N.list Opt.List
75 <*> traverse (optimize cycle) entries
77 Can.Negate expr ->
78 do func <- Names.registerGlobal ModuleName.basics N.negate
79 arg <- optimize cycle expr
80 pure $ Opt.Call func [arg]
82 Can.Binop _ home name _ left right ->
83 do optFunc <- Names.registerGlobal home name
84 optLeft <- optimize cycle left
85 optRight <- optimize cycle right
86 return (Opt.Call optFunc [optLeft, optRight])
88 Can.Lambda args body ->
89 do (argNames, destructors) <- destructArgs args
90 obody <- optimize cycle body
91 pure $ Opt.Function argNames (foldr Opt.Destruct obody destructors)
93 Can.Call func args ->
94 Opt.Call
95 <$> optimize cycle func
96 <*> traverse (optimize cycle) args
98 Can.If branches finally ->
99 let
100 optimizeBranch (condition, branch) =
101 (,)
102 <$> optimize cycle condition
103 <*> optimize cycle branch
104 in
105 Opt.If
106 <$> traverse optimizeBranch branches
107 <*> optimize cycle finally
109 Can.Let def body ->
110 optimizeDef cycle def =<< optimize cycle body
112 Can.LetRec defs body ->
113 case defs of
114 [def] ->
115 Opt.Let
116 <$> optimizePotentialTailCall cycle def
117 <*> optimize cycle body
119 _ ->
120 do obody <- optimize cycle body
121 foldM (\bod def -> optimizeDef cycle def bod) obody defs
123 Can.LetDestruct pattern expr body ->
124 do (name, destructs) <- destruct pattern
125 oexpr <- optimize cycle expr
126 obody <- optimize cycle body
127 pure $
128 Opt.Let (Opt.Def name oexpr) (foldr Opt.Destruct obody destructs)
130 Can.Case expr branches ->
131 let
132 optimizeBranch root (Can.CaseBranch pattern branch) =
133 do destructors <- destructCase root pattern
134 obranch <- optimize cycle branch
135 pure (pattern, foldr Opt.Destruct obranch destructors)
136 in
137 do temp <- Names.generate
138 oexpr <- optimize cycle expr
139 case oexpr of
140 Opt.VarLocal root ->
141 Case.optimize temp root <$> traverse (optimizeBranch root) branches
143 _ ->
144 do obranches <- traverse (optimizeBranch temp) branches
145 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
147 Can.Accessor field ->
148 Names.registerField field (Opt.Accessor field)
150 Can.Access record (A.At _ field) ->
151 do optRecord <- optimize cycle record
152 Names.registerField field (Opt.Access optRecord field)
154 Can.Update _ record updates ->
155 Names.registerFieldDict updates Opt.Update
156 <*> optimize cycle record
157 <*> traverse (optimizeUpdate cycle) updates
159 Can.Record fields ->
160 Names.registerFieldDict fields Opt.Record
161 <*> traverse (optimize cycle) fields
163 Can.Unit ->
164 Names.registerKernel N.utils Opt.Unit
166 Can.Tuple a b maybeC ->
167 Names.registerKernel N.utils Opt.Tuple
168 <*> optimize cycle a
169 <*> optimize cycle b
170 <*> traverse (optimize cycle) maybeC
172 Can.Shader _ src _ ->
173 pure (Opt.Shader src)
177-- UPDATE
180optimizeUpdate :: Cycle -> Can.FieldUpdate -> Names.Tracker Opt.Expr
181optimizeUpdate cycle (Can.FieldUpdate _ expr) =
182 optimize cycle expr
186-- DEFINITION
189optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
190optimizeDef cycle def body =
191 case def of
192 Can.Def (A.At _ name) args expr ->
193 optimizeDefHelp cycle name args expr body
195 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
196 optimizeDefHelp cycle name (map fst typedArgs) expr body
199optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
200optimizeDefHelp cycle name args expr body =
201 do oexpr <- optimize cycle expr
202 case args of
203 [] ->
204 pure $ Opt.Let (Opt.Def name oexpr) body
206 _ ->
207 do (argNames, destructors) <- destructArgs args
208 let ofunc = Opt.Function argNames (foldr Opt.Destruct oexpr destructors)
209 pure $ Opt.Let (Opt.Def name ofunc) body
213-- DESTRUCTURING
216destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [Opt.Destructor])
217destructArgs args =
218 do (argNames, destructorLists) <- unzip <$> traverse destruct args
219 return (argNames, concat destructorLists)
222destructCase :: N.Name -> Can.Pattern -> Names.Tracker [Opt.Destructor]
223destructCase rootName pattern =
224 reverse <$> destructHelp (Opt.Root rootName) pattern []
227destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Destructor])
228destruct pattern@(A.At _ ptrn) =
229 case ptrn of
230 Can.PVar name ->
231 pure (name, [])
233 Can.PAlias subPattern name ->
234 do revDs <- destructHelp (Opt.Root name) subPattern []
235 pure (name, reverse revDs)
237 _ ->
238 do name <- Names.generate
239 revDs <- destructHelp (Opt.Root name) pattern []
240 pure (name, reverse revDs)
243destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
244destructHelp path (A.At region pattern) revDs =
245 case pattern of
246 Can.PAnything ->
247 pure revDs
249 Can.PVar name ->
250 pure (Opt.Destructor name path : revDs)
252 Can.PRecord fields ->
253 let
254 toDestruct name =
255 Opt.Destructor name (Opt.Field name path)
256 in
257 Names.registerFieldList fields (map toDestruct fields ++ revDs)
259 Can.PAlias subPattern name ->
260 destructHelp (Opt.Root name) subPattern $
261 Opt.Destructor name path : revDs
263 Can.PUnit ->
264 pure revDs
266 Can.PTuple a b Nothing ->
267 destructTwo path a b revDs
269 Can.PTuple a b (Just c) ->
270 case path of
271 Opt.Root _ ->
272 destructHelp (Opt.Index Index.third path) c =<<
273 destructHelp (Opt.Index Index.second path) b =<<
274 destructHelp (Opt.Index Index.first path) a revDs
276 _ ->
277 do name <- Names.generate
278 let newRoot = Opt.Root name
279 destructHelp (Opt.Index Index.third newRoot) c =<<
280 destructHelp (Opt.Index Index.second newRoot) b =<<
281 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
283 Can.PList [] ->
284 pure revDs
286 Can.PList (hd:tl) ->
287 destructTwo path hd (A.At region (Can.PList tl)) revDs
289 Can.PCons hd tl ->
290 destructTwo path hd tl revDs
292 Can.PChr _ ->
293 pure revDs
295 Can.PStr _ ->
296 pure revDs
298 Can.PInt _ ->
299 pure revDs
301 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
302 case args of
303 [Can.PatternCtorArg _ _ arg] ->
304 case opts of
305 Can.Normal -> destructHelp (Opt.Index Index.first path) arg revDs
306 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
307 Can.Enum -> destructHelp (Opt.Index Index.first path) arg revDs
309 _ ->
310 case path of
311 Opt.Root _ ->
312 foldM (destructCtorArg path) revDs args
314 _ ->
315 do name <- Names.generate
316 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
319destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
320destructTwo path a b revDs =
321 case path of
322 Opt.Root _ ->
323 destructHelp (Opt.Index Index.second path) b =<<
324 destructHelp (Opt.Index Index.first path) a revDs
326 _ ->
327 do name <- Names.generate
328 let newRoot = Opt.Root name
329 destructHelp (Opt.Index Index.second newRoot) b =<<
330 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
333destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
334destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
335 destructHelp (Opt.Index index path) arg revDs
339-- TAIL CALL
342optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
343optimizePotentialTailCall cycle def =
344 case def of
345 Can.Def (A.At _ name) args expr ->
346 do (argNames, destructors) <- destructArgs args
347 toTailDef name argNames destructors <$>
348 optimizeTail cycle name argNames expr
350 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
351 do (argNames, destructors) <- destructArgs (map fst typedArgs)
352 toTailDef name argNames destructors <$>
353 optimizeTail cycle name argNames expr
356optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
357optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
358 case expression of
359 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
360 do oargs <- traverse (optimize cycle) args
361 if name == rootName
362 then
363 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
364 Index.LengthMatch pairs ->
365 pure $ Opt.TailCall name pairs
367 Index.LengthMismatch _ _ ->
368 do ofunc <- optimize cycle func
369 pure $ Opt.Call ofunc oargs
370 else
371 do ofunc <- optimize cycle func
372 pure $ Opt.Call ofunc oargs
374 Can.If branches finally ->
375 let
376 optimizeBranch (condition, branch) =
377 (,)
378 <$> optimize cycle condition
379 <*> optimizeTail cycle rootName argNames branch
380 in
381 Opt.If
382 <$> traverse optimizeBranch branches
383 <*> optimizeTail cycle rootName argNames finally
385 Can.Let def body ->
386 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
388 Can.LetRec defs body ->
389 case defs of
390 [def] ->
391 Opt.Let
392 <$> optimizePotentialTailCall cycle def
393 <*> optimizeTail cycle rootName argNames body
395 _ ->
396 do obody <- optimizeTail cycle rootName argNames body
397 foldM (\bod def -> optimizeDef cycle def bod) obody defs
399 Can.LetDestruct pattern expr body ->
400 do (dname, destructors) <- destruct pattern
401 oexpr <- optimize cycle expr
402 obody <- optimizeTail cycle dname argNames body
403 pure $
404 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
406 Can.Case expr branches ->
407 let
408 optimizeBranch root (Can.CaseBranch pattern branch) =
409 do destructors <- destructCase root pattern
410 obranch <- optimizeTail cycle rootName argNames branch
411 pure (pattern, foldr Opt.Destruct obranch destructors)
412 in
413 do temp <- Names.generate
414 oexpr <- optimize cycle expr
415 case oexpr of
416 Opt.VarLocal root ->
417 Case.optimize temp root <$> traverse (optimizeBranch root) branches
419 _ ->
420 do obranches <- traverse (optimizeBranch temp) branches
421 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
423 _ ->
424 optimize cycle locExpr
428-- DETECT TAIL CALLS
431toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
432toTailDef name argNames destructors body =
433 if hasTailCall body then
434 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
435 else
436 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
439hasTailCall :: Opt.Expr -> Bool
440hasTailCall expression =
441 case expression of
442 Opt.TailCall _ _ ->
443 True
445 Opt.If branches finally ->
446 hasTailCall finally || any (hasTailCall . snd) branches
448 Opt.Let _ body ->
449 hasTailCall body
451 Opt.Destruct _ body ->
452 hasTailCall body
454 Opt.Case _ _ decider jumps ->
455 decidecHasTailCall decider || any (hasTailCall . snd) jumps
457 _ ->
458 False
461decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
462decidecHasTailCall decider =
463 case decider of
464 Opt.Leaf choice ->
465 case choice of
466 Opt.Inline expr ->
467 hasTailCall expr
469 Opt.Jump _ ->
470 False
472 Opt.Chain _ success failure ->
473 decidecHasTailCall success || decidecHasTailCall failure
475 Opt.FanOut _ tests fallback ->
476 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests