on bf63cb8 help sign in
elk/src/Optimize/Expression.hs
436 lines 12 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 case args of
191 [] ->
192 do oexpr <- optimize cycle expr
193 pure $ Opt.Let (Opt.Def name oexpr) body
195 _ ->
196 do (argNames, destructors) <- destructArgs args
197 let oexpr = Opt.Function argNames (foldr Opt.Destruct body destructors)
198 pure $ Opt.Let (Opt.Def name oexpr) 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 do name <- Names.generate
255 let newRoot = Opt.Root name
256 destructHelp (Opt.Index Index.first newRoot) a =<<
257 destructHelp (Opt.Index Index.second newRoot) b =<<
258 destructHelp (Opt.Index Index.third newRoot) c (Opt.Destructor name path : revDs)
260 Can.PList [] ->
261 pure revDs
263 Can.PList (hd:tl) ->
264 destructTwo path hd (A.At region (Can.PList tl)) revDs
266 Can.PCons hd tl ->
267 destructTwo path hd tl revDs
269 Can.PChr _ ->
270 pure revDs
272 Can.PStr _ ->
273 pure revDs
275 Can.PInt _ ->
276 pure revDs
278 Can.PCtor _ _ _ _ _ args ->
279 do name <- Names.generate
280 foldM
281 (destructCtorArg (Opt.Root name))
282 (Opt.Destructor name path : revDs)
283 args
286destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
287destructTwo path a b revDs =
288 do name <- Names.generate
289 let newRoot = Opt.Root name
290 destructHelp (Opt.Index Index.first newRoot) a =<<
291 destructHelp (Opt.Index Index.second newRoot) b (Opt.Destructor name path : revDs)
294destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
295destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
296 destructHelp (Opt.Index index path) arg revDs
300-- TAIL CALL
303optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
304optimizePotentialTailCall cycle def =
305 case def of
306 Can.Def (A.At _ name) args expr ->
307 do (argNames, destructors) <- destructArgs args
308 toTailDef name argNames destructors <$>
309 optimizeTail cycle name argNames expr
311 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
312 do (argNames, destructors) <- destructArgs (map fst typedArgs)
313 toTailDef name argNames destructors <$>
314 optimizeTail cycle name argNames expr
317optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
318optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
319 case expression of
320 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
321 do oargs <- traverse (optimize cycle) args
322 if name == rootName
323 then
324 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
325 Index.LengthMatch pairs ->
326 pure $ Opt.TailCall name pairs
328 Index.LengthMismatch _ _ ->
329 do ofunc <- optimize cycle func
330 pure $ Opt.Call ofunc oargs
331 else
332 do ofunc <- optimize cycle func
333 pure $ Opt.Call ofunc oargs
335 Can.If branches finally ->
336 let
337 optimizeBranch (condition, branch) =
338 (,)
339 <$> optimize cycle condition
340 <*> optimizeTail cycle rootName argNames branch
341 in
342 Opt.If
343 <$> traverse optimizeBranch branches
344 <*> optimizeTail cycle rootName argNames finally
346 Can.Let def body ->
347 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
349 Can.LetRec defs body ->
350 case defs of
351 [def] ->
352 Opt.Let
353 <$> optimizePotentialTailCall cycle def
354 <*> optimizeTail cycle rootName argNames body
356 _ ->
357 do obody <- optimizeTail cycle rootName argNames body
358 foldM (\bod def -> optimizeDef cycle def bod) obody defs
360 Can.LetDestruct pattern expr body ->
361 do (dname, destructors) <- destruct pattern
362 oexpr <- optimize cycle expr
363 obody <- optimizeTail cycle dname argNames body
364 pure $
365 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
367 Can.Case expr branches ->
368 let
369 optimizeBranch root (Can.CaseBranch pattern branch) =
370 do destructors <- destructHelp (Opt.Root root) pattern []
371 obranch <- optimizeTail cycle rootName argNames branch
372 pure (pattern, foldr Opt.Destruct obranch destructors)
373 in
374 do temp <- Names.generate
375 oexpr <- optimize cycle expr
376 case oexpr of
377 Opt.VarLocal root ->
378 Case.optimize temp root <$> traverse (optimizeBranch root) branches
380 _ ->
381 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
383 _ ->
384 optimize cycle locExpr
388-- DETECT TAIL CALLS
391toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
392toTailDef name argNames destructors body =
393 if hasTailCall body then
394 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
395 else
396 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
399hasTailCall :: Opt.Expr -> Bool
400hasTailCall expression =
401 case expression of
402 Opt.TailCall _ _ ->
403 True
405 Opt.If branches finally ->
406 hasTailCall finally || any (hasTailCall . snd) branches
408 Opt.Let _ body ->
409 hasTailCall body
411 Opt.Destruct _ body ->
412 hasTailCall body
414 Opt.Case _ _ decider jumps ->
415 decidecHasTailCall decider || any (hasTailCall . snd) jumps
417 _ ->
418 False
421decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
422decidecHasTailCall decider =
423 case decider of
424 Opt.Leaf choice ->
425 case choice of
426 Opt.Inline expr ->
427 hasTailCall expr
429 Opt.Jump _ ->
430 False
432 Opt.Chain _ success failure ->
433 decidecHasTailCall success || decidecHasTailCall failure
435 Opt.FanOut _ tests fallback ->
436 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests