git.iko.soy/elk
on 441aeba help sign in
elk/src/Optimize.hs
550 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 var = Var.Global home name
54 deps = Set.singleton var
55 decl =
56 Env.run variantDict home $
57 ((,) name . Opt.Def) <$> optimizeExpr Nothing rawBody
58 destructors =
59 destruct (Opt.VarGlobal var) pattern
60 in
61 decl : map (second (Opt.Decl deps Nothing)) destructors
63 _ ->
64 error "bug manifesting in Optimize.optimizeDecl, please report"
68-- OPTIMIZE DEFS
71optimizeLetDef :: Can.Def -> Env.Optimizer [(Text, Opt.Def)]
72optimizeLetDef (Can.Def _ pattern@(A.A _ ptrn) expression _) =
73 let
74 (rawArgs, rawBody) =
75 Can.collectLambdas expression
76 in
77 case (ptrn, rawArgs) of
78 (P.Var name, _) ->
79 do def <- optimizeNamedDef name rawArgs rawBody
80 return [def]
82 (_, []) ->
83 do name <- Env.freshName
84 body <- optimizeExpr Nothing rawBody
85 return $ (name, Opt.Def body) : destruct (Opt.VarLocal name) pattern
87 _ ->
88 error "bug manifesting in Optimize.optimizeLetDef, please report"
91optimizeNamedDef :: Text -> [P.Canonical] -> Can.Expr -> Env.Optimizer (Text, Opt.Def)
92optimizeNamedDef name rawArgs rawBody =
93 case rawArgs of
94 [] ->
95 do hasTailCall <- Env.getTailCall
96 body <- optimizeExpr Nothing rawBody
97 Env.setTailCall hasTailCall
98 return (name, Opt.Def body)
100 _ : _ ->
101 do htc <- Env.getTailCall
102 Env.setTailCall False
104 (args, body) <- optimizeFunction (Just name) rawArgs rawBody
106 isTailCall <- Env.getTailCall
107 Env.setTailCall htc
109 return $ (,) name $
110 if isTailCall then
111 Opt.TailDef args body
112 else
113 Opt.Def (Opt.Function args body)
116destruct :: Opt.Expr -> P.Canonical -> [(Text, Opt.Def)]
117destruct root pattern =
118 map (second Opt.Def) $
119 patternToSubstitutions root pattern
123-- OPTIMIZE FUNCTION
126optimizeFunction
127 :: Maybe Text
128 -> [P.Canonical]
129 -> Can.Expr
130 -> Env.Optimizer ([Text], Opt.Expr)
131optimizeFunction maybeName patterns givenBody =
132 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
133 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
134 return (argNames, optBody)
137makeContext :: [Text] -> Text -> Context
138makeContext argNames name =
139 Just (name, length argNames, argNames)
142depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
143depatternArgs (names, rawExpr) ptrn =
144 do (name, expr) <- depattern ptrn rawExpr
145 return (name : names, expr)
149-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
152patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
153patternToSubstitutions expr (A.A _ pattern) =
154 case pattern of
155 P.Var name ->
156 [ (name, expr) ]
158 P.Literal _ ->
159 []
161 P.Record fields ->
162 map (\name -> (name, Opt.Access expr name)) fields
164 P.Anything ->
165 []
167 P.Alias alias realPattern ->
168 (alias, expr) : patternToSubstitutions expr realPattern
170 P.Ctor _ patterns ->
171 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
175-- OPTIMIZE EXPRESSIONS
178type Context = Maybe (Text, Int, [Text])
181{-| This optimizer detects tail calls. The `Context` holds a function name,
182function arity, and list of argument names. If a tail call is possible, the
183`Context` will be filled in and we check any function we run into. If we go
184down a route that cannot possibly have tail calls, we remove the context.
185-}
186optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
187optimizeExpr context annExpr@(A.A region expression) =
188 let
189 keepLooking =
190 optimizeExpr context
191 -- this *might* find a tail call, use this when something is a potential
192 -- tail call
194 justConvert =
195 optimizeExpr Nothing
196 -- this will never find tail calls, use this when you are working with
197 -- expressions that cannot be turned into tail calls
198 in
199 case expression of
200 Can.Literal lit ->
201 pure (Opt.Literal lit)
203 Can.Var var ->
204 optimizeVariable region var
206 Can.List elements ->
207 Opt.List
208 <$> traverse justConvert elements
210 Can.Binop op leftExpr rightExpr ->
211 optimizeBinop context region op leftExpr rightExpr
213 Can.Lambda _ _ ->
214 let
215 (patterns, body) =
216 Can.collectLambdas annExpr
217 in
218 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
220 Can.App _ _ ->
221 let
222 (func:args) =
223 Can.collectApps annExpr
224 in
225 case getTailCallInfo context func (length args) of
226 Just (name, argNames) ->
227 do optArgs <- traverse justConvert args
228 Env.setTailCall True
229 return (Opt.TailCall name argNames optArgs)
231 Nothing ->
232 do hasTailCall <- Env.getTailCall
233 optFunc <- justConvert func
234 optArgs <- traverse justConvert args
235 Env.setTailCall hasTailCall
236 return (Opt.Call optFunc optArgs)
238 Can.If branches finally ->
239 let
240 crawlBranch (cond,branch) =
241 (,) <$> justConvert cond <*> keepLooking branch
242 in
243 Opt.If
244 <$> traverse crawlBranch branches
245 <*> keepLooking finally
247 Can.Let defs body ->
248 do optDefs <- concat <$> traverse optimizeLetDef defs
249 Opt.Let optDefs <$> keepLooking body
251 Can.Case expr branches ->
252 do optExpr <- optimizeExpr Nothing expr
253 variantDict <- Env.getVariantDict
255 name <- Env.freshName
256 optBranches <- traverse (optimizeBranch context region name) branches
257 let optCase = Case.optimize variantDict name optBranches
258 return $ Opt.Let [ (name, Opt.Def optExpr) ] optCase
260 Can.Ctor (Var.Canonical _ name) exprs ->
261 Opt.Ctor name <$> traverse justConvert exprs
263 Can.Access record field ->
264 Opt.Access
265 <$> justConvert record
266 <*> pure field
268 Can.Update record fields ->
269 Opt.Update
270 <$> justConvert record
271 <*> traverse (traverse justConvert) fields
273 Can.Record fields ->
274 Opt.Record
275 <$> traverse (traverse justConvert) fields
277 Can.Cmd moduleName manager ->
278 Opt.Cmd moduleName <$> Env.registerEffects manager
280 Can.Sub moduleName manager ->
281 Opt.Sub moduleName <$> Env.registerEffects manager
283 Can.OutgoingPort name tipe ->
284 pure (Opt.OutgoingPort name tipe)
286 Can.IncomingPort name tipe ->
287 pure (Opt.IncomingPort name tipe)
289 Can.Program kind expr ->
290 Opt.Program kind <$> justConvert expr
292 Can.SaveEnv _ _ ->
293 error "save_the_environment should never make it to optimization phase"
295 Can.GLShader _ src _ ->
296 pure (Opt.GLShader src)
300-- OPTIMIZE VARIABLE
303optimizeVariable :: R.Region -> Var.Canonical -> Env.Optimizer Opt.Expr
304optimizeVariable region (Var.Canonical home name) =
305 case home of
306 Var.BuiltIn ->
307 pure (Opt.VarLocal name)
309 Var.Local ->
310 pure (Opt.VarLocal name)
312 Var.TopLevel modul ->
313 do let var = Var.Global modul name
314 Env.register var
315 pure (Opt.VarGlobal var)
317 Var.Module modul ->
318 do let var = Var.Global modul name
319 Env.register var
320 if name == "crash" && modul == ModuleName.inCore "Debug"
321 then
322 do here <- Env.getHome
323 pure (Opt.Crash here region Nothing)
324 else
325 pure (Opt.VarGlobal var)
329-- OPTIMIZE CASE BRANCHES
332optimizeBranch
333 :: Context
334 -> R.Region
335 -> Text
336 -> (P.Canonical, Can.Expr)
337 -> Env.Optimizer (P.Canonical, Opt.Expr)
338optimizeBranch context region exprName (rawPattern, canBranch) =
339 do optBranch <- optimizeExpr context canBranch
340 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
342 let substitutions = Map.fromList $
343 patternToSubstitutions (Opt.VarLocal exprName) pattern
345 (,) pattern <$> Inline.inline substitutions branch
348tagCrashBranch
349 :: R.Region
350 -> P.Canonical
351 -> Opt.Expr
352 -> Env.Optimizer (P.Canonical, Opt.Expr)
353tagCrashBranch region pattern@(A.A pr _) expr =
354 case expr of
355 Opt.Call (Opt.Crash home _ _) [arg] ->
356 do name <- Env.freshName
357 return
358 ( A.A pr (P.Alias name pattern)
359 , Opt.Call (Opt.Crash home region (Just (Opt.VarLocal name))) [arg]
360 )
362 _ ->
363 return (pattern, expr)
367-- DETECT TAIL CALL
370getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
371getTailCallInfo context (A.A _ function) arity =
372 case (context, function) of
373 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
375 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
376 Just (name, argNames)
378 else
379 Nothing
381 _ ->
382 Nothing
386-- DEPATTERN
387-- given a pattern and an expression, push the actual pattern matching into the
388-- expression as a case-expression if necessary. The pattern compiler will work
389-- it out from there
392depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
393depattern canPattern@(A.A _ pattern) expr =
394 let
395 ann =
396 A.A (error "the annotation added in 'depattern' should not be observed!")
398 caseExpr e branches =
399 ann (Can.Case e branches)
401 var name =
402 ann (Can.Var (Var.local name))
404 caseify =
405 do name <- Env.freshName
406 return
407 ( name
408 , caseExpr (var name) [ (canPattern, expr) ]
409 )
410 in
411 case pattern of
412 P.Var name ->
413 return (name, expr)
415 P.Literal _ ->
416 caseify
418 P.Record _ ->
419 caseify
421 P.Anything ->
422 do name <- Env.freshName
423 return (name, expr)
425 P.Alias _ _ ->
426 caseify
428 P.Ctor _ _ ->
429 caseify
433-- OPTIMIZE BINOPS
436optimizeBinop
437 :: Context
438 -> R.Region
439 -> Var.Canonical
440 -> Can.Expr
441 -> Can.Expr
442 -> Env.Optimizer Opt.Expr
443optimizeBinop context region op leftExpr rightExpr =
444 let
445 ann = A.A region
446 binop = ann (Can.Binop op leftExpr rightExpr)
447 in
448 if op == forwardApply then
450 optimizeExpr context (collect id left forwardApply binop)
452 else if op == backwardApply then
454 optimizeExpr context (collect id right backwardApply binop)
456 else if op == forwardCompose then
458 do var <- Env.freshName
459 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
460 let body = collect makeRoot left forwardCompose binop
461 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
463 else if op == backwardCompose then
465 do var <- Env.freshName
466 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
467 let body = collect makeRoot right backwardCompose binop
468 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
470 else if op == cons then
472 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
474 else
476 makeOptBinop op
477 <$> optimizeExpr Nothing leftExpr
478 <*> optimizeExpr Nothing rightExpr
481-- left-associative ((x |> f) |> g)
482forwardApply :: Var.Canonical
483forwardApply =
484 Var.inCore "Basics" "|>"
487-- right-associative (g <| (f <| x))
488backwardApply :: Var.Canonical
489backwardApply =
490 Var.inCore "Basics" "<|"
493-- left-associative ((f >> g) >> h)
494forwardCompose :: Var.Canonical
495forwardCompose =
496 Var.inCore "Basics" ">>"
499-- right-associative (h << (g << f))
500backwardCompose :: Var.Canonical
501backwardCompose =
502 Var.inCore "Basics" "<<"
505cons :: Var.Canonical
506cons =
507 Var.Canonical Var.BuiltIn "::"
510type Assoc =
511 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
514left :: Assoc
515left =
516 flip (,)
519right :: Assoc
520right =
521 (,)
524collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
525collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
526 case expr of
527 Can.Binop op leftExpr rightExpr | op == desiredOp ->
528 let
529 (func, arg) = assoc leftExpr rightExpr
530 in
531 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
533 _ ->
534 makeRoot annExpr
537makeOptBinop :: Var.Canonical -> Opt.Expr -> Opt.Expr -> Opt.Expr
538makeOptBinop (Var.Canonical home op) leftExpr rightExpr =
539 case home of
540 Var.Local ->
541 error "bug manifesting in Optimize.makeOptBinop"
543 Var.BuiltIn ->
544 error "bug manifesting in Optimize.makeOptBinop"
546 Var.Module name ->
547 Opt.Binop (Var.Global name op) leftExpr rightExpr
549 Var.TopLevel name ->
550 Opt.Binop (Var.Global name op) leftExpr rightExpr