on 57490fd help sign in
elk/src/Optimize.hs
398 lines 11 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize
4 ( optimizeExpr
5 )
6 where
9import Control.Monad (foldM)
11import qualified AST.Canonical as Can
12import qualified AST.Optimized as Opt
13import qualified AST.Module.Name as ModuleName
14import qualified Data.Index as Index
15import qualified Elm.Name as N
16import qualified Optimize.Case as Case
17import qualified Optimize.Names as Names
18import qualified Reporting.Annotation as A
22-- OPTIMIZE
26-- EXPRESSIONS
29optimizeExpr :: Can.Expr -> Names.Tracker Opt.Expr
30optimizeExpr (A.At region expression) =
31 case expression of
32 Can.VarLocal name ->
33 pure (Opt.VarLocal name)
35 Can.VarTopLevel home name ->
36 Names.registerGlobal home name
38 Can.VarKernel home name ->
39 Names.registerKernel home (Opt.VarKernel home name)
41 Can.VarForeign home name _ ->
42 Names.registerGlobal home name
44 Can.VarDebug home name _ ->
45 Names.registerKernel N.debug (Opt.VarDebug name home region Nothing)
47 Can.VarOperator _ home name _ ->
48 Names.registerGlobal home name
50 Can.Chr chr ->
51 pure (Opt.Chr chr)
53 Can.Str str ->
54 pure (Opt.Str str)
56 Can.Int int ->
57 pure (Opt.Int int)
59 Can.Float float ->
60 pure (Opt.Float float)
62 Can.List entries ->
63 Opt.List <$> traverse optimizeExpr entries
65 Can.Negate expr ->
66 do func <- Names.registerGlobal ModuleName.basics N.negate
67 arg <- optimizeExpr expr
68 pure $ Opt.Call func [arg]
70 Can.Binop _ home name _ left right ->
71 do optFunc <- Names.registerGlobal home name
72 optLeft <- optimizeExpr left
73 optRight <- optimizeExpr right
74 return (Opt.Call optFunc [optLeft, optRight])
76 Can.Lambda args body ->
77 do (argNames, defss) <- destructArgs args
78 obody <- optimizeExpr body
79 pure $ Opt.Function argNames (foldr Opt.Let obody (concat defss))
81 Can.Call func args ->
82 Opt.Call
83 <$> optimizeExpr func
84 <*> traverse optimizeExpr args
86 Can.If branches finally ->
87 let
88 optimizeBranch (condition, branch) =
89 (,)
90 <$> optimizeExpr condition
91 <*> optimizeExpr branch
92 in
93 Opt.If
94 <$> traverse optimizeBranch branches
95 <*> optimizeExpr finally
97 Can.Let def body ->
98 optimizeDef def =<< optimizeExpr body
100 Can.LetRec defs body ->
101 case defs of
102 [def] ->
103 Opt.Let
104 <$> optimizePotentialTailCall def
105 <*> optimizeExpr body
107 _ ->
108 do obody <- optimizeExpr body
109 foldM (\bod def -> optimizeDef def bod) obody defs
111 Can.LetDestruct pattern expr body ->
112 do (name, defs) <- destruct pattern
113 oexpr <- optimizeExpr expr
114 obody <- optimizeExpr body
115 pure $
116 Opt.Let (Opt.Def name oexpr) (foldr Opt.Let obody defs)
118 Can.Case expr branches ->
119 let
120 optimizeBranch root (Can.CaseBranch pattern branch) =
121 do defs <- destructHelp (Opt.VarLocal root) pattern []
122 obranch <- optimizeExpr branch
123 pure (pattern, foldr Opt.Let obranch defs)
124 in
125 do oexpr <- optimizeExpr expr
126 case oexpr of
127 Opt.VarLocal root ->
128 Case.optimize root <$> traverse (optimizeBranch root) branches
130 _ ->
131 do root <- Names.generate
132 Case.optimize root <$> traverse (optimizeBranch root) branches
134 Can.Accessor field ->
135 pure (Opt.Accessor field)
137 Can.Access record field ->
138 do optRecord <- optimizeExpr record
139 pure (Opt.Access optRecord field)
141 Can.Update record updates ->
142 Opt.Update
143 <$> optimizeExpr record
144 <*> traverse optimizeExpr updates
146 Can.Record fields ->
147 Opt.Record <$> traverse optimizeExpr fields
149 Can.Unit ->
150 Names.registerKernel N.utils Opt.Unit
152 Can.Tuple a b maybeC ->
153 Names.registerKernel N.utils Opt.Tuple
154 <*> optimizeExpr a
155 <*> optimizeExpr b
156 <*> traverse optimizeExpr maybeC
158 Can.Shader _ src _ ->
159 pure (Opt.Shader src)
163-- DEFINITION
166optimizeDef :: Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
167optimizeDef def body =
168 case def of
169 Can.Def (A.At _ name) args expr ->
170 optimizeDefHelp name args expr body
172 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
173 optimizeDefHelp name (map fst typedArgs) expr body
176optimizeDefHelp :: N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
177optimizeDefHelp name args expr body =
178 case args of
179 [] ->
180 do oexpr <- optimizeExpr expr
181 pure $ Opt.Let (Opt.Def name oexpr) body
183 _ ->
184 do (argNames, defss) <- destructArgs args
185 let func = Opt.Function argNames (foldr Opt.Let body (concat defss))
186 pure $
187 Opt.Let (Opt.Def name func) body
191-- DESTRUCTURING
194destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [[Opt.Def]])
195destructArgs args =
196 unzip <$> traverse destruct args
199destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Def])
200destruct pattern@(A.At _ ptrn) =
201 case ptrn of
202 Can.PVar name ->
203 pure (name, [])
205 Can.PAlias subPattern name ->
206 do revDefs <- destructHelp (Opt.VarLocal name) subPattern []
207 pure (name, reverse revDefs)
209 _ ->
210 do name <- Names.generate
211 revDefs <- destructHelp (Opt.VarLocal name) pattern []
212 pure (name, reverse revDefs)
215destructHelp :: Opt.Expr -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
216destructHelp root (A.At region pattern) revDefs =
217 case pattern of
218 Can.PAnything ->
219 pure revDefs
221 Can.PVar name ->
222 pure (Opt.Def name root : revDefs)
224 Can.PRecord fields ->
225 let
226 toDef field =
227 Opt.Def field (Opt.Access root field)
228 in
229 pure (map toDef fields ++ revDefs)
231 Can.PAlias subPattern name ->
232 destructHelp (Opt.VarLocal name) subPattern $
233 Opt.Def name root : revDefs
235 Can.PUnit ->
236 pure revDefs
238 Can.PTuple a b maybeC ->
239 do name <- Names.generate
240 let tuple = Opt.VarLocal name
241 let newRevDefs = Opt.Def name root : revDefs
242 case maybeC of
243 Nothing ->
244 destructHelp (Opt.CtorAccess tuple Index.second) b =<<
245 destructHelp (Opt.CtorAccess tuple Index.first) a newRevDefs
247 Just c ->
248 destructHelp (Opt.CtorAccess tuple Index.third) c =<<
249 destructHelp (Opt.CtorAccess tuple Index.second) b =<<
250 destructHelp (Opt.CtorAccess tuple Index.first) a newRevDefs
252 Can.PList [] ->
253 pure revDefs
255 Can.PList (hd:tl) ->
256 destructCons root hd (A.At region (Can.PList tl)) revDefs
258 Can.PCons hd tl ->
259 destructCons root hd tl revDefs
261 Can.PChr _ ->
262 pure revDefs
264 Can.PStr _ ->
265 pure revDefs
267 Can.PInt _ ->
268 pure revDefs
270 Can.PCtor _ _ _ _ _ args ->
271 do name <- Names.generate
272 foldM
273 (destructCtorArg (Opt.VarLocal name))
274 (Opt.Def name root : revDefs)
275 args
278destructCons :: Opt.Expr -> Can.Pattern -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
279destructCons root hd tl revDefs =
280 do name <- Names.generate
281 let list = Opt.VarLocal name
282 destructHelp (Opt.CtorAccess list Index.first) hd =<<
283 destructHelp (Opt.CtorAccess list Index.second) tl (Opt.Def name root : revDefs)
286destructCtorArg :: Opt.Expr -> [Opt.Def] -> Can.PatternCtorArg -> Names.Tracker [Opt.Def]
287destructCtorArg root revDefs (Can.PatternCtorArg index _ arg) =
288 destructHelp (Opt.CtorAccess root index) arg revDefs
292-- TAIL CALL
295optimizePotentialTailCall :: Can.Def -> Names.Tracker Opt.Def
296optimizePotentialTailCall def =
297 case def of
298 Can.Def (A.At _ name) args expr ->
299 do (argNames, defss) <- destructArgs args
300 toTailDef name argNames defss <$>
301 optimizeTailExpr name argNames expr
303 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
304 do (argNames, defss) <- destructArgs (map fst typedArgs)
305 toTailDef name argNames defss <$>
306 optimizeTailExpr name argNames expr
309optimizeTailExpr :: N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
310optimizeTailExpr name argNames locExpr@(A.At _ expression) =
311 case expression of
312 Can.Call (A.At _ (Can.VarTopLevel _ funcName)) args ->
313 if funcName == name && length args == length argNames then
314 Opt.TailCall name argNames <$> traverse optimizeExpr args
315 else
316 optimizeExpr locExpr
318 Can.If branches finally ->
319 let
320 optimizeBranch (condition, branch) =
321 (,)
322 <$> optimizeExpr condition
323 <*> optimizeTailExpr name argNames branch
324 in
325 Opt.If
326 <$> traverse optimizeBranch branches
327 <*> optimizeTailExpr name argNames finally
329 Can.Let def body ->
330 optimizeDef def =<< optimizeTailExpr name argNames body
332 Can.LetRec defs body ->
333 case defs of
334 [def] ->
335 Opt.Let
336 <$> optimizePotentialTailCall def
337 <*> optimizeTailExpr name argNames body
339 _ ->
340 do obody <- optimizeTailExpr name argNames body
341 foldM (\bod def -> optimizeDef def bod) obody defs
343 Can.LetDestruct pattern expr body ->
344 do (dname, defs) <- destruct pattern
345 oexpr <- optimizeExpr expr
346 obody <- optimizeTailExpr dname argNames body
347 pure $
348 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Let obody defs)
350 Can.Case expr branches ->
351 let
352 optimizeBranch root (Can.CaseBranch pattern branch) =
353 do defs <- destructHelp (Opt.VarLocal root) pattern []
354 obranch <- optimizeTailExpr name argNames branch
355 pure (pattern, foldr Opt.Let obranch defs)
356 in
357 do oexpr <- optimizeExpr expr
358 case oexpr of
359 Opt.VarLocal root ->
360 Case.optimize root <$> traverse (optimizeBranch root) branches
362 _ ->
363 do root <- Names.generate
364 Case.optimize root <$> traverse (optimizeBranch root) branches
366 _ ->
367 optimizeExpr locExpr
371-- DETECT TAIL CALLS
374toTailDef :: N.Name -> [N.Name] -> [[Opt.Def]] -> Opt.Expr -> Opt.Def
375toTailDef name argNames defss body =
376 if hasTailCall body then
377 Opt.TailDef name argNames (foldr Opt.Let body (concat defss))
378 else
379 Opt.Def name (Opt.Function argNames (foldr Opt.Let body (concat defss)))
382hasTailCall :: Opt.Expr -> Bool
383hasTailCall expression =
384 case expression of
385 Opt.TailCall _ _ _ ->
386 True
388 Opt.If branches finally ->
389 hasTailCall finally || any (hasTailCall . snd) branches
391 Opt.Let _ body ->
392 hasTailCall body
394 Opt.Case _ _ branches ->
395 any (hasTailCall . snd) branches
397 _ ->
398 False