on 94ebc2d help sign in
elk/src/Optimize.hs
533 lines 14 kB history raw
1module Optimize (optimize) where
3import qualified Control.Monad as M
4import qualified Data.Map as Map
5import qualified Data.Traversable as T
7import qualified AST.Expression.General as Expr
8import qualified AST.Expression.Canonical as Can
9import qualified AST.Expression.Optimized as Opt
10import qualified AST.Module as Module
11import qualified AST.Module.Name as ModuleName
12import qualified AST.Pattern as P
13import qualified AST.Variable as Var
14import qualified Optimize.Case as Case
15import qualified Optimize.DecisionTree as DT
16import qualified Optimize.Environment as Env
17import qualified Optimize.Inline as Inline
18import qualified Reporting.Annotation as A
19import qualified Reporting.Region as R
23-- OPTIMIZE
26optimize :: DT.VariantDict -> Module.Canonical -> Module.Optimized
27optimize variantDict modul@(Module.Module home _ info) =
28 let
29 (defs, _) =
30 flattenLets [] (Module.program info)
32 optDefs =
33 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
35 optInfo =
36 info { Module.program = optDefs }
37 in
38 modul { Module.info = optInfo }
41flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
42flattenLets defs annExpr@(A.A _ expr) =
43 case expr of
44 Expr.Let ds body ->
45 flattenLets (defs ++ ds) body
47 _ ->
48 (defs, annExpr)
52-- CONVERT DEFINITIONS
55optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
56optimizeDef isRoot (Can.Def (Can.Facts deps) pattern expression _) =
57 let
58 (args, canBody) =
59 Expr.collectLambdas expression
60 in
61 do home <-
62 if isRoot then
63 Just <$> Env.getHome
64 else
65 return Nothing
67 optimizeDefHelp home deps pattern args canBody
70optimizeDefHelp
71 :: Maybe ModuleName.Canonical
72 -> [Var.TopLevel]
73 -> P.Canonical
74 -> [P.Canonical]
75 -> Can.Expr
76 -> Env.Optimizer [Opt.Def]
77optimizeDefHelp home deps pattern@(A.A _ ptrn) args rawBody =
78 let
79 facts =
80 Opt.Facts home deps
81 in
82 case (ptrn, args) of
83 (P.Var name, _ : _) ->
85 do htc <- Env.getTailCall
86 Env.setTailCall False
88 (args, body) <- optimizeFunction (Just name) args rawBody
90 isTailCall <- Env.getTailCall
91 Env.setTailCall htc
93 return $
94 if isTailCall then
95 [ Opt.TailDef facts name args body ]
97 else
98 [ Opt.Def facts name (Opt.Function args body) ]
100 (P.Var name, []) ->
102 do hasTailCall <- Env.getTailCall
103 body <- optimizeExpr Nothing rawBody
104 Env.setTailCall hasTailCall
105 return [ Opt.Def facts name body ]
107 (_, []) ->
109 do name <- Env.freshName
110 optBody <- optimizeExpr Nothing rawBody
111 let optDef = Opt.Def facts name optBody
112 return (optDef : patternToDefs home name pattern)
114 _ ->
115 error "there should never be a function where the name is not a P.Var"
119-- OPTIMIZE FUNCTION
122optimizeFunction
123 :: Maybe String
124 -> [P.Canonical]
125 -> Can.Expr
126 -> Env.Optimizer ([String], Opt.Expr)
127optimizeFunction maybeName patterns givenBody =
128 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
129 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
130 return (argNames, optBody)
133makeContext :: [String] -> String -> Context
134makeContext argNames name =
135 Just (name, length argNames, argNames)
138depatternArgs :: ([String], Can.Expr) -> P.Canonical -> Env.Optimizer ([String], Can.Expr)
139depatternArgs (names, rawExpr) ptrn =
140 do (name, expr) <- depattern ptrn rawExpr
141 return (name : names, expr)
144patternToDefs :: Maybe ModuleName.Canonical -> String -> P.Canonical -> [Opt.Def]
145patternToDefs home rootName pattern =
146 let
147 (deps, rootExpr) =
148 case home of
149 Nothing ->
150 ( []
151 , Opt.Var (Var.local rootName)
152 )
154 Just moduleName ->
155 ( [ Var.TopLevelVar moduleName rootName ]
156 , Opt.Var (Var.topLevel moduleName rootName)
157 )
159 toDef (name, accessExpr) =
160 Opt.Def (Opt.Facts home deps) name accessExpr
161 in
162 map toDef (patternToSubstitutions rootExpr pattern)
166-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
169patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(String, Opt.Expr)]
170patternToSubstitutions expr (A.A _ pattern) =
171 case pattern of
172 P.Var name ->
173 [ (name, expr) ]
175 P.Literal _ ->
176 []
178 P.Record fields ->
179 map (\name -> (name, Opt.Access expr name)) fields
181 P.Anything ->
182 []
184 P.Alias alias realPattern ->
185 (alias, expr) : patternToSubstitutions expr realPattern
187 P.Data _ patterns ->
188 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
192-- OPTIMIZE EXPRESSIONS
195type Context = Maybe (String, Int, [String])
198{-| This optimizer detects tail calls. The `Context` holds a function name,
199function arity, and list of argument names. If a tail call is possible, the
200`Context` will be filled in and we check any function we run into. If we go
201down a route that cannot possibly have tail calls, we remove the context.
202-}
203optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
204optimizeExpr context annExpr@(A.A region expression) =
205 let
206 keepLooking =
207 optimizeExpr context
208 -- this *might* find a tail call, use this when something is a potential
209 -- tail call
211 justConvert =
212 optimizeExpr Nothing
213 -- this will never find tail calls, use this when you are working with
214 -- expressions that cannot be turned into tail calls
215 in
216 case expression of
217 Expr.Literal lit ->
218 pure (Opt.Literal lit)
220 Expr.Var name ->
221 if name == Var.inCore ["Debug"] "crash" then
222 do home <- Env.getHome
223 pure (Opt.Crash home region Nothing)
225 else
226 pure (Opt.Var name)
228 Expr.Range lowExpr highExpr ->
229 Opt.Range
230 <$> justConvert lowExpr
231 <*> justConvert highExpr
233 Expr.ExplicitList elements ->
234 Opt.ExplicitList
235 <$> T.traverse justConvert elements
237 Expr.Binop op leftExpr rightExpr ->
238 optimizeBinop context region op leftExpr rightExpr
240 Expr.Lambda _ _ ->
241 let
242 (patterns, body) =
243 Expr.collectLambdas annExpr
244 in
245 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
247 Expr.App _ _ ->
248 let
249 (func:args) =
250 Expr.collectApps annExpr
251 in
252 case isTailCall context func (length args) of
253 Just (name, argNames) ->
254 do optArgs <- T.traverse justConvert args
255 Env.setTailCall True
256 return (Opt.TailCall name argNames optArgs)
258 Nothing ->
259 do hasTailCall <- Env.getTailCall
260 optFunc <- justConvert func
261 optArgs <- T.traverse justConvert args
262 Env.setTailCall hasTailCall
263 return (Opt.Call optFunc optArgs)
265 Expr.If branches finally ->
266 let
267 crawlBranch (cond,branch) =
268 (,) <$> justConvert cond <*> keepLooking branch
269 in
270 Opt.If
271 <$> T.traverse crawlBranch branches
272 <*> keepLooking finally
274 Expr.Let defs body ->
275 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
276 Opt.Let optDefs <$> keepLooking body
278 Expr.Case expr branches ->
279 do optExpr <- optimizeExpr Nothing expr
280 variantDict <- Env.getVariantDict
282 name <- Env.freshName
283 optBranches <- T.traverse (optimizeBranch context region name) branches
284 let optCase = Case.optimize variantDict name optBranches
285 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
287 Expr.Data name exprs ->
288 Opt.Data name <$> T.traverse justConvert exprs
290 Expr.Access record field ->
291 Opt.Access
292 <$> justConvert record
293 <*> pure field
295 Expr.Update record fields ->
296 Opt.Update
297 <$> justConvert record
298 <*> T.traverse (mapSnd justConvert) fields
300 Expr.Record fields ->
301 Opt.Record
302 <$> T.traverse (mapSnd justConvert) fields
304 Expr.Cmd moduleName ->
305 pure (Opt.Cmd moduleName)
307 Expr.Sub moduleName ->
308 pure (Opt.Sub moduleName)
310 Expr.ForeignCmd name tipe ->
311 pure (Opt.ForeignCmd name tipe)
313 Expr.ForeignSub name tipe ->
314 pure (Opt.ForeignSub name tipe)
316 Expr.SaveEnv _ _ ->
317 error "save_the_environment should never make it to optimization phase"
319 Expr.GLShader uid src gltipe ->
320 pure (Opt.GLShader uid src gltipe)
323mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
324mapSnd func (x, a) =
325 (,) x <$> func a
329-- OPTIMIZE CASE BRANCHES
332optimizeBranch
333 :: Context
334 -> R.Region
335 -> String
336 -> (P.Canonical, Can.Expr)
337 -> Env.Optimizer (P.Canonical, Opt.Expr)
338optimizeBranch context region exprName (rawPattern, canBranch) =
339 let
340 root =
341 Opt.Var (Var.Canonical Var.Local exprName)
342 in
343 do optBranch <- optimizeExpr context canBranch
344 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
346 let substitutions =
347 Map.fromList (patternToSubstitutions root pattern)
349 (,) pattern <$> Inline.inline substitutions branch
352tagCrashBranch
353 :: R.Region
354 -> P.Canonical
355 -> Opt.Expr
356 -> Env.Optimizer (P.Canonical, Opt.Expr)
357tagCrashBranch region pattern@(A.A pr _) expr =
358 case expr of
359 Opt.Call (Opt.Crash home _ _) [arg] ->
360 do name <- Env.freshName
361 return
362 ( A.A pr (P.Alias name pattern)
363 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
364 )
366 _ ->
367 return (pattern, expr)
371-- DETECT TAIL CALL
374isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
375isTailCall context (A.A _ function) arity =
376 case (context, function) of
377 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
379 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
380 Just (name, argNames)
382 else
383 Nothing
385 _ ->
386 Nothing
390-- DEPATTERN
391-- given a pattern and an expression, push the actual pattern matching into the
392-- expression as a case-expression if necessary. The pattern compiler will work
393-- it out from there
396depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (String, Can.Expr)
397depattern canPattern@(A.A _ pattern) expr =
398 let
399 ann =
400 A.A (error "the annotation added in 'depattern' should not be observed!")
402 caseExpr e branches =
403 ann (Expr.Case e branches)
405 var name =
406 ann (Expr.Var (Var.local name))
408 caseify =
409 do name <- Env.freshName
410 return
411 ( name
412 , caseExpr (var name) [ (canPattern, expr) ]
413 )
414 in
415 case pattern of
416 P.Var name ->
417 return (name, expr)
419 P.Literal _ ->
420 caseify
422 P.Record _ ->
423 caseify
425 P.Anything ->
426 do name <- Env.freshName
427 return (name, expr)
429 P.Alias _ _ ->
430 caseify
432 P.Data _ _ ->
433 caseify
437-- OPTIMIZE BINOPS
440optimizeBinop
441 :: Context
442 -> R.Region
443 -> Var.Canonical
444 -> Can.Expr
445 -> Can.Expr
446 -> Env.Optimizer Opt.Expr
447optimizeBinop context region op leftExpr rightExpr =
448 let
449 ann = A.A region
450 binop = ann (Expr.Binop op leftExpr rightExpr)
451 in
452 if op == forwardApply then
454 optimizeExpr context (collect id left forwardApply binop)
456 else if op == backwardApply then
458 optimizeExpr context (collect id right backwardApply binop)
460 else if op == forwardCompose then
462 do var <- Env.freshName
463 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
464 let body = collect makeRoot left forwardCompose binop
465 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
467 else if op == backwardCompose then
469 do var <- Env.freshName
470 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
471 let body = collect makeRoot right backwardCompose binop
472 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
474 else if op == Var.Canonical Var.BuiltIn "::" then
476 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
478 else
480 Opt.Binop op
481 <$> optimizeExpr Nothing leftExpr
482 <*> optimizeExpr Nothing rightExpr
485-- left-associative ((x |> f) |> g)
486forwardApply :: Var.Canonical
487forwardApply =
488 Var.inCore ["Basics"] "|>"
491-- right-associative (g <| (f <| x))
492backwardApply :: Var.Canonical
493backwardApply =
494 Var.inCore ["Basics"] "<|"
497-- left-associative ((f >> g) >> h)
498forwardCompose :: Var.Canonical
499forwardCompose =
500 Var.inCore ["Basics"] ">>"
503-- right-associative (h << (g << f))
504backwardCompose :: Var.Canonical
505backwardCompose =
506 Var.inCore ["Basics"] "<<"
509type Assoc =
510 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
513left :: Assoc
514left =
515 flip (,)
518right :: Assoc
519right =
520 (,)
523collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
524collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
525 case expr of
526 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
527 let
528 (func, arg) = assoc leftExpr rightExpr
529 in
530 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
532 _ ->
533 makeRoot annExpr