on 2297bbf help sign in
elk/src/Optimize.hs
463 lines 12 kB history raw
1module Optimize (optimize) where
3import Control.Applicative
4import qualified Control.Monad as M
5import qualified Data.Map as Map
6import qualified Data.Traversable as T
8import qualified AST.Expression.General as Expr
9import qualified AST.Expression.Canonical as Can
10import qualified AST.Expression.Optimized as Opt
11import qualified AST.Module as Module
12import qualified AST.Module.Name as ModuleName
13import qualified AST.Pattern as P
14import qualified AST.Variable as Var
15import Elm.Utils ((|>))
16import qualified Optimize.Environment as Env
17import qualified Optimize.Patterns as Patterns
18import qualified Reporting.Annotation as A
19import qualified Reporting.Region as R
22-- OPTIMIZE - turn a canonical module into a list of optimized defs
24optimize :: Module.CanonicalModule -> Module.Optimized
25optimize (Module.Module home _ _ _ _ body) =
26 let
27 (defs, _) =
28 flattenLets [] (Module.program body)
30 variantInfo (_, tags) =
31 let
32 variants = length tags
33 in
34 map (\(tag, _) -> (Var.fromModule home tag, variants)) tags
36 variantDict =
37 Module.datatypes body
38 |> Map.elems
39 |> concatMap variantInfo
40 |> Map.fromList
41 |> Map.insert (Var.builtin "_Tuple2") 1
42 in
43 Env.run variantDict (concat <$> mapM (optimizeDef (Just home)) defs)
46flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
47flattenLets defs annExpr@(A.A _ expr) =
48 case expr of
49 Expr.Let ds body ->
50 flattenLets (defs ++ ds) body
52 _ ->
53 (defs, annExpr)
56-- CONVERT DEFINITIONS
58optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Env.Optimizer [Opt.Def]
59optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
60 let
61 (args, rawBody) =
62 Expr.collectLambdas expression
63 in
64 case (ptrn, args) of
65 (P.Var name, _ : _) ->
67 do htc <- Env.getTailCall
68 Env.setTailCall False
70 (args, body) <- optimizeFunction (Just name) args rawBody
72 isTailCall <- Env.getTailCall
73 Env.setTailCall htc
75 let facts = Opt.Facts home deps
77 return $
78 if isTailCall then
80 [ Opt.TailDef facts name args body ]
82 else
84 [ Opt.Def facts name (Opt.Function args body) ]
86 (P.Var name, []) ->
88 do hasTailCall <- Env.getTailCall
89 body <- optimizeExpr Nothing rawBody
90 Env.setTailCall hasTailCall
91 return [ Opt.Def (Opt.Facts home deps) name body ]
93 (_, []) ->
95 do name <- Env.freshName
96 optBody <- optimizeExpr Nothing rawBody
97 let optDef = Opt.Def (Opt.Facts home deps) name optBody
98 return (optDef : patternToDefs home name pattern)
100 _ ->
101 error "there should never be a function where the name is not a P.Var"
104-- OPTIMIZE FUNCTION
106optimizeFunction
107 :: Maybe String
108 -> [P.CanonicalPattern]
109 -> Can.Expr
110 -> Env.Optimizer ([String], Opt.Expr)
111optimizeFunction maybeName patterns givenBody =
112 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
113 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
114 return (argNames, optBody)
117makeContext :: [String] -> String -> Context
118makeContext argNames name =
119 Just (name, length argNames, argNames)
122depatternArgs
123 :: ([String], Can.Expr)
124 -> P.CanonicalPattern
125 -> Env.Optimizer ([String], Can.Expr)
126depatternArgs (names, rawExpr) ptrn =
127 do (name, expr) <- depattern ptrn rawExpr
128 return (name : names, expr)
131patternToDefs
132 :: Maybe ModuleName.Canonical
133 -> String
134 -> P.CanonicalPattern
135 -> [Opt.Def]
136patternToDefs home rootName pattern =
137 let
138 (deps, rootExpr) =
139 case home of
140 Nothing ->
141 ( []
142 , Opt.Var (Var.local rootName)
143 )
145 Just moduleName ->
146 ( [ Var.TopLevelVar moduleName rootName ]
147 , Opt.Var (Var.topLevel moduleName rootName)
148 )
150 toDef (name, accessExpr) =
151 Opt.Def (Opt.Facts home deps) name accessExpr
152 in
153 map toDef (substitutions rootExpr pattern)
156substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
157substitutions expr (A.A _ pattern) =
158 case pattern of
159 P.Var name ->
160 [ (name, expr) ]
162 P.Literal _ ->
163 []
165 P.Record fields ->
166 map (\name -> (name, Opt.Access expr name)) fields
168 P.Anything ->
169 []
171 P.Alias alias realPattern ->
172 (alias, expr) : substitutions expr realPattern
174 P.Data _ patterns ->
175 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
178-- OPTIMIZE EXPRESSIONS
180type Context = Maybe (String, Int, [String])
183{-| This optimizer detects tail calls. The `Context` holds a function name,
184function arity, and list of argument names. If a tail call is possible, the
185`Context` will be filled in and we check any function we run into. If we go
186down a route that cannot possibly have tail calls, we remove the context.
187-}
188optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
189optimizeExpr context annExpr@(A.A region expression) =
190 let
191 keepLooking =
192 optimizeExpr context
193 -- this *might* find a tail call, use this when something is a potential
194 -- tail call
196 justConvert =
197 optimizeExpr Nothing
198 -- this will never find tail calls, use this when you are working with
199 -- expressions that cannot be turned into tail calls
200 in
201 case expression of
202 Expr.Literal lit ->
203 pure (Opt.Literal lit)
205 Expr.Var name ->
206 if name == Var.inCore ["Debug"] "crash" then
207 pure (Opt.Crash region Nothing)
209 else
210 pure (Opt.Var name)
212 Expr.Range lowExpr highExpr ->
213 Opt.Range
214 <$> justConvert lowExpr
215 <*> justConvert highExpr
217 Expr.ExplicitList elements ->
218 Opt.ExplicitList
219 <$> T.traverse justConvert elements
221 Expr.Binop op leftExpr rightExpr ->
222 optimizeBinop context region op leftExpr rightExpr
224 Expr.Lambda _ _ ->
225 let
226 (patterns, body) =
227 Expr.collectLambdas annExpr
228 in
229 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
231 Expr.App _ _ ->
232 let
233 (func:args) =
234 Expr.collectApps annExpr
235 in
236 case isTailCall context func (length args) of
237 Just (name, argNames) ->
238 do optArgs <- T.traverse justConvert args
239 Env.setTailCall True
240 return (Opt.TailCall name argNames optArgs)
242 Nothing ->
243 do hasTailCall <- Env.getTailCall
244 optFunc <- justConvert func
245 optArgs <- T.traverse justConvert args
246 Env.setTailCall hasTailCall
247 return (Opt.Call optFunc optArgs)
249 Expr.If branches finally ->
250 let
251 crawlBranch (cond,branch) =
252 (,) <$> justConvert cond <*> keepLooking branch
253 in
254 Opt.If
255 <$> T.traverse crawlBranch branches
256 <*> keepLooking finally
258 Expr.Let defs body ->
259 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
260 Opt.Let optDefs <$> keepLooking body
262 Expr.Case expr branches ->
263 do optExpr <- optimizeExpr Nothing expr
264 optBranches <- T.traverse (mapSnd (optimizeExpr context)) branches
265 variantDict <- Env.getVariantDict
266 Patterns.optimize variantDict region optExpr optBranches
268 Expr.Data name exprs ->
269 Opt.Data name <$> T.traverse justConvert exprs
271 Expr.Access record field ->
272 Opt.Access
273 <$> justConvert record
274 <*> pure field
276 Expr.Update record fields ->
277 Opt.Update
278 <$> justConvert record
279 <*> T.traverse (mapSnd justConvert) fields
281 Expr.Record fields ->
282 Opt.Record
283 <$> T.traverse (mapSnd justConvert) fields
285 Expr.GLShader uid src gltipe ->
286 pure (Opt.GLShader uid src gltipe)
288 Expr.Port impl ->
289 Opt.Port <$>
290 case impl of
291 Expr.In name tipe ->
292 pure (Expr.In name tipe)
294 Expr.Out name expr tipe ->
295 (\e -> Expr.Out name e tipe) <$> justConvert expr
297 Expr.Task name expr tipe ->
298 (\e -> Expr.Task name e tipe) <$> justConvert expr
301mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
302mapSnd func (x, a) =
303 (,) x <$> func a
306-- DETECT TAIL CALL
308isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
309isTailCall context (A.A _ function) arity =
310 case (context, function) of
311 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
313 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
314 Just (name, argNames)
316 else
317 Nothing
319 _ ->
320 Nothing
323-- DEPATTERN
324-- given a pattern and an expression, push the actual pattern matching into the
325-- expression as a case-expression if necessary. The pattern compiler will work
326-- it out from there
328depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
329depattern canPattern@(A.A _ pattern) expr =
330 let
331 ann =
332 A.A (error "the annotation added in 'depattern' should not be observed!")
334 caseExpr e branches =
335 ann (Expr.Case e branches)
337 var name =
338 ann (Expr.Var (Var.local name))
340 caseify =
341 do name <- Env.freshName
342 return
343 ( name
344 , caseExpr (var name) [ (canPattern, expr) ]
345 )
346 in
347 case pattern of
348 P.Var name ->
349 return (name, expr)
351 P.Literal _ ->
352 caseify
354 P.Record _ ->
355 caseify
357 P.Anything ->
358 do name <- Env.freshName
359 return (name, expr)
361 P.Alias _ _ ->
362 caseify
364 P.Data _ _ ->
365 caseify
368-- OPTIMIZE BINOPS
370optimizeBinop
371 :: Context
372 -> R.Region
373 -> Var.Canonical
374 -> Can.Expr
375 -> Can.Expr
376 -> Env.Optimizer Opt.Expr
377optimizeBinop context region op leftExpr rightExpr =
378 let
379 ann = A.A region
380 binop = ann (Expr.Binop op leftExpr rightExpr)
381 in
382 if op == forwardApply then
384 optimizeExpr context (collect id left forwardApply binop)
386 else if op == backwardApply then
388 optimizeExpr context (collect id right backwardApply binop)
390 else if op == forwardCompose then
392 do var <- Env.freshName
393 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
394 let body = collect makeRoot left forwardCompose binop
395 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
397 else if op == backwardCompose then
399 do var <- Env.freshName
400 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
401 let body = collect makeRoot right backwardCompose binop
402 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
404 else if op == Var.Canonical Var.BuiltIn "::" then
406 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
408 else
410 Opt.Binop op
411 <$> optimizeExpr Nothing leftExpr
412 <*> optimizeExpr Nothing rightExpr
415-- left-associative ((x |> f) |> g)
416forwardApply :: Var.Canonical
417forwardApply =
418 Var.inCore ["Basics"] "|>"
421-- right-associative (g <| (f <| x))
422backwardApply :: Var.Canonical
423backwardApply =
424 Var.inCore ["Basics"] "<|"
427-- left-associative ((f >> g) >> h)
428forwardCompose :: Var.Canonical
429forwardCompose =
430 Var.inCore ["Basics"] ">>"
433-- right-associative (h << (g << f))
434backwardCompose :: Var.Canonical
435backwardCompose =
436 Var.inCore ["Basics"] "<<"
439type Assoc =
440 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
443left :: Assoc
444left =
445 flip (,)
448right :: Assoc
449right =
450 (,)
453collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
454collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
455 case expr of
456 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
457 let
458 (func, arg) = assoc leftExpr rightExpr
459 in
460 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
462 _ ->
463 makeRoot annExpr