git.iko.soy/elk
on fe1e5d2 help sign in
elk/src/Optimize.hs
524 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 module_@(Module.Module _ home _ info) =
28 let
29 (Module.Canonical docs exports imports body) =
30 info
32 (defs, _) =
33 flattenLets [] (Module.program body)
35 optDefs =
36 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
38 optBody =
39 body { Module.program = optDefs }
40 in
41 module_
42 { Module.info =
43 Module.Optimized docs exports imports optBody
44 }
47flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
48flattenLets defs annExpr@(A.A _ expr) =
49 case expr of
50 Expr.Let ds body ->
51 flattenLets (defs ++ ds) body
53 _ ->
54 (defs, annExpr)
58-- CONVERT DEFINITIONS
61optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
62optimizeDef isRoot (Can.Definition (Can.Facts deps) pattern expression _) =
63 let
64 (args, canBody) =
65 Expr.collectLambdas expression
66 in
67 do home <-
68 if isRoot then
69 Just <$> Env.getHome
70 else
71 return Nothing
73 optimizeDefHelp home deps pattern args canBody
76optimizeDefHelp
77 :: Maybe ModuleName.Canonical
78 -> [Var.TopLevel]
79 -> P.Canonical
80 -> [P.Canonical]
81 -> Can.Expr
82 -> Env.Optimizer [Opt.Def]
83optimizeDefHelp home deps pattern@(A.A _ ptrn) args rawBody =
84 let
85 facts =
86 Opt.Facts home deps
87 in
88 case (ptrn, args) of
89 (P.Var name, _ : _) ->
91 do htc <- Env.getTailCall
92 Env.setTailCall False
94 (args, body) <- optimizeFunction (Just name) args rawBody
96 isTailCall <- Env.getTailCall
97 Env.setTailCall htc
99 return $
100 if isTailCall then
101 [ Opt.TailDef facts name args body ]
103 else
104 [ Opt.Def facts name (Opt.Function args body) ]
106 (P.Var name, []) ->
108 do hasTailCall <- Env.getTailCall
109 body <- optimizeExpr Nothing rawBody
110 Env.setTailCall hasTailCall
111 return [ Opt.Def facts name body ]
113 (_, []) ->
115 do name <- Env.freshName
116 optBody <- optimizeExpr Nothing rawBody
117 let optDef = Opt.Def facts name optBody
118 return (optDef : patternToDefs home name pattern)
120 _ ->
121 error "there should never be a function where the name is not a P.Var"
125-- OPTIMIZE FUNCTION
128optimizeFunction
129 :: Maybe String
130 -> [P.Canonical]
131 -> Can.Expr
132 -> Env.Optimizer ([String], Opt.Expr)
133optimizeFunction maybeName patterns givenBody =
134 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
135 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
136 return (argNames, optBody)
139makeContext :: [String] -> String -> Context
140makeContext argNames name =
141 Just (name, length argNames, argNames)
144depatternArgs :: ([String], Can.Expr) -> P.Canonical -> Env.Optimizer ([String], Can.Expr)
145depatternArgs (names, rawExpr) ptrn =
146 do (name, expr) <- depattern ptrn rawExpr
147 return (name : names, expr)
150patternToDefs :: Maybe ModuleName.Canonical -> String -> P.Canonical -> [Opt.Def]
151patternToDefs home rootName pattern =
152 let
153 (deps, rootExpr) =
154 case home of
155 Nothing ->
156 ( []
157 , Opt.Var (Var.local rootName)
158 )
160 Just moduleName ->
161 ( [ Var.TopLevelVar moduleName rootName ]
162 , Opt.Var (Var.topLevel moduleName rootName)
163 )
165 toDef (name, accessExpr) =
166 Opt.Def (Opt.Facts home deps) name accessExpr
167 in
168 map toDef (patternToSubstitutions rootExpr pattern)
172-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
175patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(String, Opt.Expr)]
176patternToSubstitutions expr (A.A _ pattern) =
177 case pattern of
178 P.Var name ->
179 [ (name, expr) ]
181 P.Literal _ ->
182 []
184 P.Record fields ->
185 map (\name -> (name, Opt.Access expr name)) fields
187 P.Anything ->
188 []
190 P.Alias alias realPattern ->
191 (alias, expr) : patternToSubstitutions expr realPattern
193 P.Data _ patterns ->
194 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
198-- OPTIMIZE EXPRESSIONS
201type Context = Maybe (String, Int, [String])
204{-| This optimizer detects tail calls. The `Context` holds a function name,
205function arity, and list of argument names. If a tail call is possible, the
206`Context` will be filled in and we check any function we run into. If we go
207down a route that cannot possibly have tail calls, we remove the context.
208-}
209optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
210optimizeExpr context annExpr@(A.A region expression) =
211 let
212 keepLooking =
213 optimizeExpr context
214 -- this *might* find a tail call, use this when something is a potential
215 -- tail call
217 justConvert =
218 optimizeExpr Nothing
219 -- this will never find tail calls, use this when you are working with
220 -- expressions that cannot be turned into tail calls
221 in
222 case expression of
223 Expr.Literal lit ->
224 pure (Opt.Literal lit)
226 Expr.Var name ->
227 if name == Var.inCore ["Debug"] "crash" then
228 do home <- Env.getHome
229 pure (Opt.Crash home region Nothing)
231 else
232 pure (Opt.Var name)
234 Expr.Range lowExpr highExpr ->
235 Opt.Range
236 <$> justConvert lowExpr
237 <*> justConvert highExpr
239 Expr.ExplicitList elements ->
240 Opt.ExplicitList
241 <$> T.traverse justConvert elements
243 Expr.Binop op leftExpr rightExpr ->
244 optimizeBinop context region op leftExpr rightExpr
246 Expr.Lambda _ _ ->
247 let
248 (patterns, body) =
249 Expr.collectLambdas annExpr
250 in
251 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
253 Expr.App _ _ ->
254 let
255 (func:args) =
256 Expr.collectApps annExpr
257 in
258 case isTailCall context func (length args) of
259 Just (name, argNames) ->
260 do optArgs <- T.traverse justConvert args
261 Env.setTailCall True
262 return (Opt.TailCall name argNames optArgs)
264 Nothing ->
265 do hasTailCall <- Env.getTailCall
266 optFunc <- justConvert func
267 optArgs <- T.traverse justConvert args
268 Env.setTailCall hasTailCall
269 return (Opt.Call optFunc optArgs)
271 Expr.If branches finally ->
272 let
273 crawlBranch (cond,branch) =
274 (,) <$> justConvert cond <*> keepLooking branch
275 in
276 Opt.If
277 <$> T.traverse crawlBranch branches
278 <*> keepLooking finally
280 Expr.Let defs body ->
281 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
282 Opt.Let optDefs <$> keepLooking body
284 Expr.Case expr branches ->
285 do optExpr <- optimizeExpr Nothing expr
286 variantDict <- Env.getVariantDict
288 name <- Env.freshName
289 optBranches <- T.traverse (optimizeBranch context region name) branches
290 let optCase = Case.optimize variantDict name optBranches
291 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
293 Expr.Data name exprs ->
294 Opt.Data name <$> T.traverse justConvert exprs
296 Expr.Access record field ->
297 Opt.Access
298 <$> justConvert record
299 <*> pure field
301 Expr.Update record fields ->
302 Opt.Update
303 <$> justConvert record
304 <*> T.traverse (mapSnd justConvert) fields
306 Expr.Record fields ->
307 Opt.Record
308 <$> T.traverse (mapSnd justConvert) fields
310 Expr.GLShader uid src gltipe ->
311 pure (Opt.GLShader uid src gltipe)
314mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
315mapSnd func (x, a) =
316 (,) x <$> func a
320-- OPTIMIZE CASE BRANCHES
323optimizeBranch
324 :: Context
325 -> R.Region
326 -> String
327 -> (P.Canonical, Can.Expr)
328 -> Env.Optimizer (P.Canonical, Opt.Expr)
329optimizeBranch context region exprName (rawPattern, canBranch) =
330 let
331 root =
332 Opt.Var (Var.Canonical Var.Local exprName)
333 in
334 do optBranch <- optimizeExpr context canBranch
335 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
337 let substitutions =
338 Map.fromList (patternToSubstitutions root pattern)
340 (,) pattern <$> Inline.inline substitutions branch
343tagCrashBranch
344 :: R.Region
345 -> P.Canonical
346 -> Opt.Expr
347 -> Env.Optimizer (P.Canonical, Opt.Expr)
348tagCrashBranch region pattern@(A.A pr _) expr =
349 case expr of
350 Opt.Call (Opt.Crash home _ _) [arg] ->
351 do name <- Env.freshName
352 return
353 ( A.A pr (P.Alias name pattern)
354 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
355 )
357 _ ->
358 return (pattern, expr)
362-- DETECT TAIL CALL
365isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
366isTailCall context (A.A _ function) arity =
367 case (context, function) of
368 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
370 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
371 Just (name, argNames)
373 else
374 Nothing
376 _ ->
377 Nothing
381-- DEPATTERN
382-- given a pattern and an expression, push the actual pattern matching into the
383-- expression as a case-expression if necessary. The pattern compiler will work
384-- it out from there
387depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (String, Can.Expr)
388depattern canPattern@(A.A _ pattern) expr =
389 let
390 ann =
391 A.A (error "the annotation added in 'depattern' should not be observed!")
393 caseExpr e branches =
394 ann (Expr.Case e branches)
396 var name =
397 ann (Expr.Var (Var.local name))
399 caseify =
400 do name <- Env.freshName
401 return
402 ( name
403 , caseExpr (var name) [ (canPattern, expr) ]
404 )
405 in
406 case pattern of
407 P.Var name ->
408 return (name, expr)
410 P.Literal _ ->
411 caseify
413 P.Record _ ->
414 caseify
416 P.Anything ->
417 do name <- Env.freshName
418 return (name, expr)
420 P.Alias _ _ ->
421 caseify
423 P.Data _ _ ->
424 caseify
428-- OPTIMIZE BINOPS
431optimizeBinop
432 :: Context
433 -> R.Region
434 -> Var.Canonical
435 -> Can.Expr
436 -> Can.Expr
437 -> Env.Optimizer Opt.Expr
438optimizeBinop context region op leftExpr rightExpr =
439 let
440 ann = A.A region
441 binop = ann (Expr.Binop op leftExpr rightExpr)
442 in
443 if op == forwardApply then
445 optimizeExpr context (collect id left forwardApply binop)
447 else if op == backwardApply then
449 optimizeExpr context (collect id right backwardApply binop)
451 else if op == forwardCompose then
453 do var <- Env.freshName
454 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
455 let body = collect makeRoot left forwardCompose binop
456 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
458 else if op == backwardCompose then
460 do var <- Env.freshName
461 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
462 let body = collect makeRoot right backwardCompose binop
463 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
465 else if op == Var.Canonical Var.BuiltIn "::" then
467 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
469 else
471 Opt.Binop op
472 <$> optimizeExpr Nothing leftExpr
473 <*> optimizeExpr Nothing rightExpr
476-- left-associative ((x |> f) |> g)
477forwardApply :: Var.Canonical
478forwardApply =
479 Var.inCore ["Basics"] "|>"
482-- right-associative (g <| (f <| x))
483backwardApply :: Var.Canonical
484backwardApply =
485 Var.inCore ["Basics"] "<|"
488-- left-associative ((f >> g) >> h)
489forwardCompose :: Var.Canonical
490forwardCompose =
491 Var.inCore ["Basics"] ">>"
494-- right-associative (h << (g << f))
495backwardCompose :: Var.Canonical
496backwardCompose =
497 Var.inCore ["Basics"] "<<"
500type Assoc =
501 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
504left :: Assoc
505left =
506 flip (,)
509right :: Assoc
510right =
511 (,)
514collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
515collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
516 case expr of
517 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
518 let
519 (func, arg) = assoc leftExpr rightExpr
520 in
521 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
523 _ ->
524 makeRoot annExpr