git.iko.soy/elk
on 247359c help sign in
elk/src/Optimize/Expression.hs
412 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.VarDebug home name _ ->
53 Names.registerKernel N.debug (Opt.VarDebug name home region Nothing)
55 Can.VarOperator _ home name _ ->
56 Names.registerGlobal home name
58 Can.Chr chr ->
59 Names.registerKernel N.utils (Opt.Chr chr)
61 Can.Str str ->
62 pure (Opt.Str str)
64 Can.Int int ->
65 pure (Opt.Int int)
67 Can.Float float ->
68 pure (Opt.Float float)
70 Can.List entries ->
71 Opt.List <$> traverse (optimize cycle) entries
73 Can.Negate expr ->
74 do func <- Names.registerGlobal ModuleName.basics N.negate
75 arg <- optimize cycle expr
76 pure $ Opt.Call func [arg]
78 Can.Binop _ home name _ left right ->
79 do optFunc <- Names.registerGlobal home name
80 optLeft <- optimize cycle left
81 optRight <- optimize cycle right
82 return (Opt.Call optFunc [optLeft, optRight])
84 Can.Lambda args body ->
85 do (argNames, defss) <- destructArgs args
86 obody <- optimize cycle body
87 pure $ Opt.Function argNames (foldr Opt.Let obody (concat defss))
89 Can.Call func args ->
90 Opt.Call
91 <$> optimize cycle func
92 <*> traverse (optimize cycle) args
94 Can.If branches finally ->
95 let
96 optimizeBranch (condition, branch) =
97 (,)
98 <$> optimize cycle condition
99 <*> optimize cycle branch
100 in
101 Opt.If
102 <$> traverse optimizeBranch branches
103 <*> optimize cycle finally
105 Can.Let def body ->
106 optimizeDef cycle def =<< optimize cycle body
108 Can.LetRec defs body ->
109 case defs of
110 [def] ->
111 Opt.Let
112 <$> optimizePotentialTailCall cycle def
113 <*> optimize cycle body
115 _ ->
116 do obody <- optimize cycle body
117 foldM (\bod def -> optimizeDef cycle def bod) obody defs
119 Can.LetDestruct pattern expr body ->
120 do (name, defs) <- destruct pattern
121 oexpr <- optimize cycle expr
122 obody <- optimize cycle body
123 pure $
124 Opt.Let (Opt.Def name oexpr) (foldr Opt.Let obody defs)
126 Can.Case expr branches ->
127 let
128 optimizeBranch root (Can.CaseBranch pattern branch) =
129 do defs <- destructHelp (Opt.VarLocal root) pattern []
130 obranch <- optimize cycle branch
131 pure (pattern, foldr Opt.Let obranch defs)
132 in
133 do temp <- Names.generate
134 oexpr <- optimize cycle expr
135 case oexpr of
136 Opt.VarLocal root ->
137 Case.optimize temp root <$> traverse (optimizeBranch root) branches
139 _ ->
140 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
142 Can.Accessor field ->
143 Names.registerField field (Opt.Accessor field)
145 Can.Access record field ->
146 do optRecord <- optimize cycle record
147 Names.registerField field (Opt.Access optRecord field)
149 Can.Update record updates ->
150 Names.registerFieldDict updates Opt.Update
151 <*> optimize cycle record
152 <*> traverse (optimize cycle) updates
154 Can.Record fields ->
155 Names.registerFieldDict fields Opt.Record
156 <*> traverse (optimize cycle) fields
158 Can.Unit ->
159 Names.registerKernel N.utils Opt.Unit
161 Can.Tuple a b maybeC ->
162 Names.registerKernel N.utils Opt.Tuple
163 <*> optimize cycle a
164 <*> optimize cycle b
165 <*> traverse (optimize cycle) maybeC
167 Can.Shader _ src _ ->
168 pure (Opt.Shader src)
172-- DEFINITION
175optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
176optimizeDef cycle def body =
177 case def of
178 Can.Def (A.At _ name) args expr ->
179 optimizeDefHelp cycle name args expr body
181 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
182 optimizeDefHelp cycle name (map fst typedArgs) expr body
185optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
186optimizeDefHelp cycle name args expr body =
187 case args of
188 [] ->
189 do oexpr <- optimize cycle expr
190 pure $ Opt.Let (Opt.Def name oexpr) body
192 _ ->
193 do (argNames, defss) <- destructArgs args
194 let oexpr = Opt.Function argNames (foldr Opt.Let body (concat defss))
195 pure $ Opt.Let (Opt.Def name oexpr) body
199-- DESTRUCTURING
202destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [[Opt.Def]])
203destructArgs args =
204 unzip <$> traverse destruct args
207destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Def])
208destruct pattern@(A.At _ ptrn) =
209 case ptrn of
210 Can.PVar name ->
211 pure (name, [])
213 Can.PAlias subPattern name ->
214 do revDefs <- destructHelp (Opt.VarLocal name) subPattern []
215 pure (name, reverse revDefs)
217 _ ->
218 do name <- Names.generate
219 revDefs <- destructHelp (Opt.VarLocal name) pattern []
220 pure (name, reverse revDefs)
223destructHelp :: Opt.Expr -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
224destructHelp root (A.At region pattern) revDefs =
225 case pattern of
226 Can.PAnything ->
227 pure revDefs
229 Can.PVar name ->
230 pure (Opt.Def name root : revDefs)
232 Can.PRecord fields ->
233 let toDef name = Opt.Def name (Opt.Access root name) in
234 Names.registerFieldList fields (map toDef fields ++ revDefs)
236 Can.PAlias subPattern name ->
237 destructHelp (Opt.VarLocal name) subPattern $
238 Opt.Def name root : revDefs
240 Can.PUnit ->
241 pure revDefs
243 Can.PTuple a b maybeC ->
244 do name <- Names.generate
245 let tuple = Opt.VarLocal name
246 let newRevDefs = Opt.Def name root : revDefs
247 case maybeC of
248 Nothing ->
249 destructHelp (Opt.Index tuple 1) b =<<
250 destructHelp (Opt.Index tuple 0) a newRevDefs
252 Just c ->
253 destructHelp (Opt.Index tuple 2) c =<<
254 destructHelp (Opt.Index tuple 1) b =<<
255 destructHelp (Opt.Index tuple 0) a newRevDefs
257 Can.PList [] ->
258 pure revDefs
260 Can.PList (hd:tl) ->
261 destructCons root hd (A.At region (Can.PList tl)) revDefs
263 Can.PCons hd tl ->
264 destructCons root hd tl revDefs
266 Can.PChr _ ->
267 pure revDefs
269 Can.PStr _ ->
270 pure revDefs
272 Can.PInt _ ->
273 pure revDefs
275 Can.PCtor _ _ _ _ _ _ args ->
276 do name <- Names.generate
277 foldM
278 (destructCtorArg (Opt.VarLocal name))
279 (Opt.Def name root : revDefs)
280 args
283destructCons :: Opt.Expr -> Can.Pattern -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
284destructCons root hd tl revDefs =
285 do name <- Names.generate
286 let list = Opt.VarLocal name
287 destructHelp (Opt.Index list 0) hd =<<
288 destructHelp (Opt.Index list 1) tl (Opt.Def name root : revDefs)
291destructCtorArg :: Opt.Expr -> [Opt.Def] -> Can.PatternCtorArg -> Names.Tracker [Opt.Def]
292destructCtorArg root revDefs (Can.PatternCtorArg index _ arg) =
293 destructHelp (Opt.Index root (Index.toZeroBased index)) arg revDefs
297-- TAIL CALL
300optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
301optimizePotentialTailCall cycle def =
302 case def of
303 Can.Def (A.At _ name) args expr ->
304 do (argNames, defss) <- destructArgs args
305 toTailDef name argNames defss <$>
306 optimizeTail cycle name argNames expr
308 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
309 do (argNames, defss) <- destructArgs (map fst typedArgs)
310 toTailDef name argNames defss <$>
311 optimizeTail cycle name argNames expr
314optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
315optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
316 case expression of
317 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
318 do oargs <- traverse (optimize cycle) args
319 if name == rootName
320 then
321 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
322 Index.LengthMatch pairs ->
323 pure $ Opt.TailCall name pairs
325 Index.LengthMismatch _ _ ->
326 do ofunc <- optimize cycle func
327 pure $ Opt.Call ofunc oargs
328 else
329 do ofunc <- optimize cycle func
330 pure $ Opt.Call ofunc oargs
332 Can.If branches finally ->
333 let
334 optimizeBranch (condition, branch) =
335 (,)
336 <$> optimize cycle condition
337 <*> optimizeTail cycle rootName argNames branch
338 in
339 Opt.If
340 <$> traverse optimizeBranch branches
341 <*> optimizeTail cycle rootName argNames finally
343 Can.Let def body ->
344 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
346 Can.LetRec defs body ->
347 case defs of
348 [def] ->
349 Opt.Let
350 <$> optimizePotentialTailCall cycle def
351 <*> optimizeTail cycle rootName argNames body
353 _ ->
354 do obody <- optimizeTail cycle rootName argNames body
355 foldM (\bod def -> optimizeDef cycle def bod) obody defs
357 Can.LetDestruct pattern expr body ->
358 do (dname, defs) <- destruct pattern
359 oexpr <- optimize cycle expr
360 obody <- optimizeTail cycle dname argNames body
361 pure $
362 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Let obody defs)
364 Can.Case expr branches ->
365 let
366 optimizeBranch root (Can.CaseBranch pattern branch) =
367 do defs <- destructHelp (Opt.VarLocal root) pattern []
368 obranch <- optimizeTail cycle rootName argNames branch
369 pure (pattern, foldr Opt.Let obranch defs)
370 in
371 do temp <- Names.generate
372 oexpr <- optimize cycle expr
373 case oexpr of
374 Opt.VarLocal root ->
375 Case.optimize temp root <$> traverse (optimizeBranch root) branches
377 _ ->
378 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
380 _ ->
381 optimize cycle locExpr
385-- DETECT TAIL CALLS
388toTailDef :: N.Name -> [N.Name] -> [[Opt.Def]] -> Opt.Expr -> Opt.Def
389toTailDef name argNames defss body =
390 if hasTailCall body then
391 Opt.TailDef name argNames (foldr Opt.Let body (concat defss))
392 else
393 Opt.Def name (Opt.Function argNames (foldr Opt.Let body (concat defss)))
396hasTailCall :: Opt.Expr -> Bool
397hasTailCall expression =
398 case expression of
399 Opt.TailCall _ _ ->
400 True
402 Opt.If branches finally ->
403 hasTailCall finally || any (hasTailCall . snd) branches
405 Opt.Let _ body ->
406 hasTailCall body
408 Opt.Case _ _ _ branches ->
409 any (hasTailCall . snd) branches
411 _ ->
412 False