git.iko.soy/elk
on fe7bbbd help sign in
elk/src/Optimize.hs
547 lines 14 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize (optimize) where
5import Control.Arrow (second)
6import qualified Control.Monad as M
7import Data.Monoid ((<>))
8import qualified Data.Map as Map
9import qualified Data.Set as Set
10import qualified Data.Text as Text
11import Data.Text (Text)
13import qualified AST.Expression.Canonical as Can
14import qualified AST.Expression.Optimized as Opt
15import qualified AST.Module.Name as ModuleName
16import qualified AST.Pattern as P
17import qualified AST.Variable as Var
18import qualified Optimize.Case as Case
19import qualified Optimize.DecisionTree as DT
20import qualified Optimize.Environment as Env
21import qualified Optimize.Inline as Inline
22import qualified Reporting.Annotation as A
23import qualified Reporting.Region as R
27-- OPTIMIZE
30optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [(Text, Opt.Decl)]
31optimize variantDict home defs =
32 concat $ zipWith (optimizeDecl variantDict home) [1 .. length defs] defs
36-- OPTIMIZE DECLS
39optimizeDecl :: DT.VariantDict -> ModuleName.Canonical -> Int -> Can.Def -> [(Text, Opt.Decl)]
40optimizeDecl variantDict home index (Can.Def _ pattern@(A.A _ ptrn) expression _) =
41 let
42 (rawArgs, rawBody) =
43 Can.collectLambdas expression
44 in
45 case (ptrn, rawArgs) of
46 (P.Var name, _) ->
47 (:[]) $ Env.run variantDict home $
48 optimizeNamedDef name rawArgs rawBody
50 (_, []) ->
51 let
52 name = "_" <> Text.pack (show index)
53 deps = Set.singleton (Var.Global home name)
54 decl =
55 Env.run variantDict home $
56 ((,) name . Opt.Def) <$> optimizeExpr Nothing rawBody
57 destructors =
58 destruct (Opt.VarGlobal home name) pattern
59 in
60 decl : map (second (Opt.Decl deps Nothing)) destructors
62 _ ->
63 error "bug manifesting in Optimize.optimizeDecl, please report"
67-- OPTIMIZE DEFS
70optimizeLetDef :: Can.Def -> Env.Optimizer [(Text, Opt.Def)]
71optimizeLetDef (Can.Def _ pattern@(A.A _ ptrn) expression _) =
72 let
73 (rawArgs, rawBody) =
74 Can.collectLambdas expression
75 in
76 case (ptrn, rawArgs) of
77 (P.Var name, _) ->
78 do def <- optimizeNamedDef name rawArgs rawBody
79 return [def]
81 (_, []) ->
82 do name <- Env.freshName
83 body <- optimizeExpr Nothing rawBody
84 return $ (name, Opt.Def body) : destruct (Opt.VarLocal name) pattern
86 _ ->
87 error "bug manifesting in Optimize.optimizeLetDef, please report"
90optimizeNamedDef :: Text -> [P.Canonical] -> Can.Expr -> Env.Optimizer (Text, Opt.Def)
91optimizeNamedDef name rawArgs rawBody =
92 case rawArgs of
93 [] ->
94 do hasTailCall <- Env.getTailCall
95 body <- optimizeExpr Nothing rawBody
96 Env.setTailCall hasTailCall
97 return (name, Opt.Def body)
99 _ : _ ->
100 do htc <- Env.getTailCall
101 Env.setTailCall False
103 (args, body) <- optimizeFunction (Just name) rawArgs rawBody
105 isTailCall <- Env.getTailCall
106 Env.setTailCall htc
108 return $ (,) name $
109 if isTailCall then
110 Opt.TailDef args body
111 else
112 Opt.Def (Opt.Function args body)
115destruct :: Opt.Expr -> P.Canonical -> [(Text, Opt.Def)]
116destruct root pattern =
117 map (second Opt.Def) $
118 patternToSubstitutions root pattern
122-- OPTIMIZE FUNCTION
125optimizeFunction
126 :: Maybe Text
127 -> [P.Canonical]
128 -> Can.Expr
129 -> Env.Optimizer ([Text], Opt.Expr)
130optimizeFunction maybeName patterns givenBody =
131 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
132 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
133 return (argNames, optBody)
136makeContext :: [Text] -> Text -> Context
137makeContext argNames name =
138 Just (name, length argNames, argNames)
141depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
142depatternArgs (names, rawExpr) ptrn =
143 do (name, expr) <- depattern ptrn rawExpr
144 return (name : names, expr)
148-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
151patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
152patternToSubstitutions expr (A.A _ pattern) =
153 case pattern of
154 P.Var name ->
155 [ (name, expr) ]
157 P.Literal _ ->
158 []
160 P.Record fields ->
161 map (\name -> (name, Opt.Access expr name)) fields
163 P.Anything ->
164 []
166 P.Alias alias realPattern ->
167 (alias, expr) : patternToSubstitutions expr realPattern
169 P.Ctor _ patterns ->
170 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
174-- OPTIMIZE EXPRESSIONS
177type Context = Maybe (Text, Int, [Text])
180{-| This optimizer detects tail calls. The `Context` holds a function name,
181function arity, and list of argument names. If a tail call is possible, the
182`Context` will be filled in and we check any function we run into. If we go
183down a route that cannot possibly have tail calls, we remove the context.
184-}
185optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
186optimizeExpr context annExpr@(A.A region expression) =
187 let
188 keepLooking =
189 optimizeExpr context
190 -- this *might* find a tail call, use this when something is a potential
191 -- tail call
193 justConvert =
194 optimizeExpr Nothing
195 -- this will never find tail calls, use this when you are working with
196 -- expressions that cannot be turned into tail calls
197 in
198 case expression of
199 Can.Literal lit ->
200 pure (Opt.Literal lit)
202 Can.Var var ->
203 optimizeVariable region var
205 Can.List elements ->
206 Opt.List
207 <$> traverse justConvert elements
209 Can.Binop op leftExpr rightExpr ->
210 optimizeBinop context region op leftExpr rightExpr
212 Can.Lambda _ _ ->
213 let
214 (patterns, body) =
215 Can.collectLambdas annExpr
216 in
217 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
219 Can.App _ _ ->
220 let
221 (func:args) =
222 Can.collectApps annExpr
223 in
224 case getTailCallInfo context func (length args) of
225 Just (name, argNames) ->
226 do optArgs <- traverse justConvert args
227 Env.setTailCall True
228 return (Opt.TailCall name argNames optArgs)
230 Nothing ->
231 do hasTailCall <- Env.getTailCall
232 optFunc <- justConvert func
233 optArgs <- traverse justConvert args
234 Env.setTailCall hasTailCall
235 return (Opt.Call optFunc optArgs)
237 Can.If branches finally ->
238 let
239 crawlBranch (cond,branch) =
240 (,) <$> justConvert cond <*> keepLooking branch
241 in
242 Opt.If
243 <$> traverse crawlBranch branches
244 <*> keepLooking finally
246 Can.Let defs body ->
247 do optDefs <- concat <$> traverse optimizeLetDef defs
248 Opt.Let optDefs <$> keepLooking body
250 Can.Case expr branches ->
251 do optExpr <- optimizeExpr Nothing expr
252 variantDict <- Env.getVariantDict
254 name <- Env.freshName
255 optBranches <- traverse (optimizeBranch context region name) branches
256 let optCase = Case.optimize variantDict name optBranches
257 return $ Opt.Let [ (name, Opt.Def optExpr) ] optCase
259 Can.Ctor (Var.Canonical _ name) exprs ->
260 Opt.Ctor name <$> traverse justConvert exprs
262 Can.Access record field ->
263 Opt.Access
264 <$> justConvert record
265 <*> pure field
267 Can.Update record fields ->
268 Opt.Update
269 <$> justConvert record
270 <*> traverse (traverse justConvert) fields
272 Can.Record fields ->
273 Opt.Record
274 <$> traverse (traverse justConvert) fields
276 Can.Cmd moduleName manager ->
277 Opt.Cmd moduleName <$> Env.registerEffects manager
279 Can.Sub moduleName manager ->
280 Opt.Sub moduleName <$> Env.registerEffects manager
282 Can.OutgoingPort name tipe ->
283 pure (Opt.OutgoingPort name tipe)
285 Can.IncomingPort name tipe ->
286 pure (Opt.IncomingPort name tipe)
288 Can.Program kind expr ->
289 Opt.Program kind <$> justConvert expr
291 Can.SaveEnv _ _ ->
292 error "save_the_environment should never make it to optimization phase"
294 Can.GLShader _ src _ ->
295 pure (Opt.GLShader src)
299-- OPTIMIZE VARIABLE
302optimizeVariable :: R.Region -> Var.Canonical -> Env.Optimizer Opt.Expr
303optimizeVariable region (Var.Canonical home name) =
304 case home of
305 Var.BuiltIn ->
306 pure (Opt.VarLocal name)
308 Var.Local ->
309 pure (Opt.VarLocal name)
311 Var.TopLevel modul ->
312 do Env.register (Var.Global modul name)
313 pure (Opt.VarGlobal modul name)
315 Var.Module modul ->
316 do Env.register (Var.Global modul name)
317 if name == "crash" && modul == ModuleName.inCore "Debug"
318 then
319 do here <- Env.getHome
320 pure (Opt.Crash here region Nothing)
321 else
322 pure (Opt.VarGlobal modul name)
326-- OPTIMIZE CASE BRANCHES
329optimizeBranch
330 :: Context
331 -> R.Region
332 -> Text
333 -> (P.Canonical, Can.Expr)
334 -> Env.Optimizer (P.Canonical, Opt.Expr)
335optimizeBranch context region exprName (rawPattern, canBranch) =
336 do optBranch <- optimizeExpr context canBranch
337 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
339 let substitutions = Map.fromList $
340 patternToSubstitutions (Opt.VarLocal exprName) pattern
342 (,) pattern <$> Inline.inline substitutions branch
345tagCrashBranch
346 :: R.Region
347 -> P.Canonical
348 -> Opt.Expr
349 -> Env.Optimizer (P.Canonical, Opt.Expr)
350tagCrashBranch region pattern@(A.A pr _) expr =
351 case expr of
352 Opt.Call (Opt.Crash home _ _) [arg] ->
353 do name <- Env.freshName
354 return
355 ( A.A pr (P.Alias name pattern)
356 , Opt.Call (Opt.Crash home region (Just (Opt.VarLocal name))) [arg]
357 )
359 _ ->
360 return (pattern, expr)
364-- DETECT TAIL CALL
367getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
368getTailCallInfo context (A.A _ function) arity =
369 case (context, function) of
370 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
372 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
373 Just (name, argNames)
375 else
376 Nothing
378 _ ->
379 Nothing
383-- DEPATTERN
384-- given a pattern and an expression, push the actual pattern matching into the
385-- expression as a case-expression if necessary. The pattern compiler will work
386-- it out from there
389depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
390depattern canPattern@(A.A _ pattern) expr =
391 let
392 ann =
393 A.A (error "the annotation added in 'depattern' should not be observed!")
395 caseExpr e branches =
396 ann (Can.Case e branches)
398 var name =
399 ann (Can.Var (Var.local name))
401 caseify =
402 do name <- Env.freshName
403 return
404 ( name
405 , caseExpr (var name) [ (canPattern, expr) ]
406 )
407 in
408 case pattern of
409 P.Var name ->
410 return (name, expr)
412 P.Literal _ ->
413 caseify
415 P.Record _ ->
416 caseify
418 P.Anything ->
419 do name <- Env.freshName
420 return (name, expr)
422 P.Alias _ _ ->
423 caseify
425 P.Ctor _ _ ->
426 caseify
430-- OPTIMIZE BINOPS
433optimizeBinop
434 :: Context
435 -> R.Region
436 -> Var.Canonical
437 -> Can.Expr
438 -> Can.Expr
439 -> Env.Optimizer Opt.Expr
440optimizeBinop context region op leftExpr rightExpr =
441 let
442 ann = A.A region
443 binop = ann (Can.Binop op leftExpr rightExpr)
444 in
445 if op == forwardApply then
447 optimizeExpr context (collect id left forwardApply binop)
449 else if op == backwardApply then
451 optimizeExpr context (collect id right backwardApply binop)
453 else if op == forwardCompose then
455 do var <- Env.freshName
456 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
457 let body = collect makeRoot left forwardCompose binop
458 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
460 else if op == backwardCompose then
462 do var <- Env.freshName
463 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
464 let body = collect makeRoot right backwardCompose binop
465 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
467 else if op == cons then
469 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
471 else
473 makeOptBinop op
474 <$> optimizeExpr Nothing leftExpr
475 <*> optimizeExpr Nothing rightExpr
478-- left-associative ((x |> f) |> g)
479forwardApply :: Var.Canonical
480forwardApply =
481 Var.inCore "Basics" "|>"
484-- right-associative (g <| (f <| x))
485backwardApply :: Var.Canonical
486backwardApply =
487 Var.inCore "Basics" "<|"
490-- left-associative ((f >> g) >> h)
491forwardCompose :: Var.Canonical
492forwardCompose =
493 Var.inCore "Basics" ">>"
496-- right-associative (h << (g << f))
497backwardCompose :: Var.Canonical
498backwardCompose =
499 Var.inCore "Basics" "<<"
502cons :: Var.Canonical
503cons =
504 Var.Canonical Var.BuiltIn "::"
507type Assoc =
508 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
511left :: Assoc
512left =
513 flip (,)
516right :: Assoc
517right =
518 (,)
521collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
522collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
523 case expr of
524 Can.Binop op leftExpr rightExpr | op == desiredOp ->
525 let
526 (func, arg) = assoc leftExpr rightExpr
527 in
528 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
530 _ ->
531 makeRoot annExpr
534makeOptBinop :: Var.Canonical -> Opt.Expr -> Opt.Expr -> Opt.Expr
535makeOptBinop (Var.Canonical home op) leftExpr rightExpr =
536 case home of
537 Var.Local ->
538 error "bug manifesting in Optimize.makeOptBinop"
540 Var.BuiltIn ->
541 error "bug manifesting in Optimize.makeOptBinop"
543 Var.Module name ->
544 Opt.Binop name op leftExpr rightExpr
546 Var.TopLevel name ->
547 Opt.Binop name op leftExpr rightExpr