on 684504e help sign in
elk/src/Optimize.hs
517 lines 13 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.Name as ModuleName
11import qualified AST.Pattern as P
12import qualified AST.Variable as Var
13import qualified Optimize.Case as Case
14import qualified Optimize.DecisionTree as DT
15import qualified Optimize.Environment as Env
16import qualified Optimize.Inline as Inline
17import qualified Reporting.Annotation as A
18import qualified Reporting.Region as R
22-- OPTIMIZE
25optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [Opt.Def]
26optimize variantDict home defs =
27 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
31-- CONVERT DEFINITIONS
34optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
35optimizeDef isRoot (Can.Def (Can.Facts deps) pattern expression _) =
36 let
37 (args, canBody) =
38 Expr.collectLambdas expression
40 maybeGetHome =
41 if isRoot then
42 Just <$> Env.getHome
44 else
45 return Nothing
46 in
47 do home <- maybeGetHome
48 optimizeDefHelp home deps pattern args canBody
51optimizeDefHelp
52 :: Maybe ModuleName.Canonical
53 -> [Var.TopLevel]
54 -> P.Canonical
55 -> [P.Canonical]
56 -> Can.Expr
57 -> Env.Optimizer [Opt.Def]
58optimizeDefHelp home deps pattern@(A.A _ ptrn) args rawBody =
59 let
60 facts =
61 Opt.Facts home deps
62 in
63 case (ptrn, args) of
64 (P.Var name, _ : _) ->
66 do htc <- Env.getTailCall
67 Env.setTailCall False
69 (args, body) <- optimizeFunction (Just name) args rawBody
71 isTailCall <- Env.getTailCall
72 Env.setTailCall htc
74 return $
75 if isTailCall then
76 [ Opt.TailDef facts name args body ]
78 else
79 [ Opt.Def facts name (Opt.Function args body) ]
81 (P.Var name, []) ->
83 do hasTailCall <- Env.getTailCall
84 body <- optimizeExpr Nothing rawBody
85 Env.setTailCall hasTailCall
86 return [ Opt.Def facts name body ]
88 (_, []) ->
90 do name <- Env.freshName
91 optBody <- optimizeExpr Nothing rawBody
92 let optDef = Opt.Def facts name optBody
93 return (optDef : patternToDefs home name pattern)
95 _ ->
96 error "there should never be a function where the name is not a P.Var"
100-- OPTIMIZE FUNCTION
103optimizeFunction
104 :: Maybe String
105 -> [P.Canonical]
106 -> Can.Expr
107 -> Env.Optimizer ([String], Opt.Expr)
108optimizeFunction maybeName patterns givenBody =
109 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
110 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
111 return (argNames, optBody)
114makeContext :: [String] -> String -> Context
115makeContext argNames name =
116 Just (name, length argNames, argNames)
119depatternArgs :: ([String], Can.Expr) -> P.Canonical -> Env.Optimizer ([String], Can.Expr)
120depatternArgs (names, rawExpr) ptrn =
121 do (name, expr) <- depattern ptrn rawExpr
122 return (name : names, expr)
125patternToDefs :: Maybe ModuleName.Canonical -> String -> P.Canonical -> [Opt.Def]
126patternToDefs home rootName pattern =
127 let
128 (deps, rootExpr) =
129 case home of
130 Nothing ->
131 ( []
132 , Opt.Var (Var.local rootName)
133 )
135 Just moduleName ->
136 ( [ Var.TopLevelVar moduleName rootName ]
137 , Opt.Var (Var.topLevel moduleName rootName)
138 )
140 toDef (name, accessExpr) =
141 Opt.Def (Opt.Facts home deps) name accessExpr
142 in
143 map toDef (patternToSubstitutions rootExpr pattern)
147-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
150patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(String, Opt.Expr)]
151patternToSubstitutions expr (A.A _ pattern) =
152 case pattern of
153 P.Var name ->
154 [ (name, expr) ]
156 P.Literal _ ->
157 []
159 P.Record fields ->
160 map (\name -> (name, Opt.Access expr name)) fields
162 P.Anything ->
163 []
165 P.Alias alias realPattern ->
166 (alias, expr) : patternToSubstitutions expr realPattern
168 P.Data _ patterns ->
169 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
173-- OPTIMIZE EXPRESSIONS
176type Context = Maybe (String, Int, [String])
179{-| This optimizer detects tail calls. The `Context` holds a function name,
180function arity, and list of argument names. If a tail call is possible, the
181`Context` will be filled in and we check any function we run into. If we go
182down a route that cannot possibly have tail calls, we remove the context.
183-}
184optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
185optimizeExpr context annExpr@(A.A region expression) =
186 let
187 keepLooking =
188 optimizeExpr context
189 -- this *might* find a tail call, use this when something is a potential
190 -- tail call
192 justConvert =
193 optimizeExpr Nothing
194 -- this will never find tail calls, use this when you are working with
195 -- expressions that cannot be turned into tail calls
196 in
197 case expression of
198 Expr.Literal lit ->
199 pure (Opt.Literal lit)
201 Expr.Var name ->
202 if name == Var.inCore ["Debug"] "crash" then
203 do home <- Env.getHome
204 pure (Opt.Crash home region Nothing)
206 else
207 pure (Opt.Var name)
209 Expr.Range lowExpr highExpr ->
210 Opt.Range
211 <$> justConvert lowExpr
212 <*> justConvert highExpr
214 Expr.ExplicitList elements ->
215 Opt.ExplicitList
216 <$> T.traverse justConvert elements
218 Expr.Binop op leftExpr rightExpr ->
219 optimizeBinop context region op leftExpr rightExpr
221 Expr.Lambda _ _ ->
222 let
223 (patterns, body) =
224 Expr.collectLambdas annExpr
225 in
226 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
228 Expr.App _ _ ->
229 let
230 (func:args) =
231 Expr.collectApps annExpr
232 in
233 case isTailCall context func (length args) of
234 Just (name, argNames) ->
235 do optArgs <- T.traverse justConvert args
236 Env.setTailCall True
237 return (Opt.TailCall name argNames optArgs)
239 Nothing ->
240 do hasTailCall <- Env.getTailCall
241 optFunc <- justConvert func
242 optArgs <- T.traverse justConvert args
243 Env.setTailCall hasTailCall
244 return (Opt.Call optFunc optArgs)
246 Expr.If branches finally ->
247 let
248 crawlBranch (cond,branch) =
249 (,) <$> justConvert cond <*> keepLooking branch
250 in
251 Opt.If
252 <$> T.traverse crawlBranch branches
253 <*> keepLooking finally
255 Expr.Let defs body ->
256 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
257 Opt.Let optDefs <$> keepLooking body
259 Expr.Case expr branches ->
260 do optExpr <- optimizeExpr Nothing expr
261 variantDict <- Env.getVariantDict
263 name <- Env.freshName
264 optBranches <- T.traverse (optimizeBranch context region name) branches
265 let optCase = Case.optimize variantDict name optBranches
266 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
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.Cmd moduleName ->
286 pure (Opt.Cmd moduleName)
288 Expr.Sub moduleName ->
289 pure (Opt.Sub moduleName)
291 Expr.ForeignCmd name tipe ->
292 pure (Opt.ForeignCmd name tipe)
294 Expr.ForeignSub name tipe ->
295 pure (Opt.ForeignSub name tipe)
297 Expr.Program kind expr ->
298 Opt.Program kind <$> justConvert expr
300 Expr.SaveEnv _ _ ->
301 error "save_the_environment should never make it to optimization phase"
303 Expr.GLShader uid src gltipe ->
304 pure (Opt.GLShader uid src gltipe)
307mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
308mapSnd func (x, a) =
309 (,) x <$> func a
313-- OPTIMIZE CASE BRANCHES
316optimizeBranch
317 :: Context
318 -> R.Region
319 -> String
320 -> (P.Canonical, Can.Expr)
321 -> Env.Optimizer (P.Canonical, Opt.Expr)
322optimizeBranch context region exprName (rawPattern, canBranch) =
323 let
324 root =
325 Opt.Var (Var.Canonical Var.Local exprName)
326 in
327 do optBranch <- optimizeExpr context canBranch
328 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
330 let substitutions =
331 Map.fromList (patternToSubstitutions root pattern)
333 (,) pattern <$> Inline.inline substitutions branch
336tagCrashBranch
337 :: R.Region
338 -> P.Canonical
339 -> Opt.Expr
340 -> Env.Optimizer (P.Canonical, Opt.Expr)
341tagCrashBranch region pattern@(A.A pr _) expr =
342 case expr of
343 Opt.Call (Opt.Crash home _ _) [arg] ->
344 do name <- Env.freshName
345 return
346 ( A.A pr (P.Alias name pattern)
347 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
348 )
350 _ ->
351 return (pattern, expr)
355-- DETECT TAIL CALL
358isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
359isTailCall context (A.A _ function) arity =
360 case (context, function) of
361 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
363 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
364 Just (name, argNames)
366 else
367 Nothing
369 _ ->
370 Nothing
374-- DEPATTERN
375-- given a pattern and an expression, push the actual pattern matching into the
376-- expression as a case-expression if necessary. The pattern compiler will work
377-- it out from there
380depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (String, Can.Expr)
381depattern canPattern@(A.A _ pattern) expr =
382 let
383 ann =
384 A.A (error "the annotation added in 'depattern' should not be observed!")
386 caseExpr e branches =
387 ann (Expr.Case e branches)
389 var name =
390 ann (Expr.Var (Var.local name))
392 caseify =
393 do name <- Env.freshName
394 return
395 ( name
396 , caseExpr (var name) [ (canPattern, expr) ]
397 )
398 in
399 case pattern of
400 P.Var name ->
401 return (name, expr)
403 P.Literal _ ->
404 caseify
406 P.Record _ ->
407 caseify
409 P.Anything ->
410 do name <- Env.freshName
411 return (name, expr)
413 P.Alias _ _ ->
414 caseify
416 P.Data _ _ ->
417 caseify
421-- OPTIMIZE BINOPS
424optimizeBinop
425 :: Context
426 -> R.Region
427 -> Var.Canonical
428 -> Can.Expr
429 -> Can.Expr
430 -> Env.Optimizer Opt.Expr
431optimizeBinop context region op leftExpr rightExpr =
432 let
433 ann = A.A region
434 binop = ann (Expr.Binop op leftExpr rightExpr)
435 in
436 if op == forwardApply then
438 optimizeExpr context (collect id left forwardApply binop)
440 else if op == backwardApply then
442 optimizeExpr context (collect id right backwardApply binop)
444 else if op == forwardCompose then
446 do var <- Env.freshName
447 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
448 let body = collect makeRoot left forwardCompose binop
449 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
451 else if op == backwardCompose 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 right backwardCompose binop
456 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
458 else if op == Var.Canonical Var.BuiltIn "::" then
460 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
462 else
464 Opt.Binop op
465 <$> optimizeExpr Nothing leftExpr
466 <*> optimizeExpr Nothing rightExpr
469-- left-associative ((x |> f) |> g)
470forwardApply :: Var.Canonical
471forwardApply =
472 Var.inCore ["Basics"] "|>"
475-- right-associative (g <| (f <| x))
476backwardApply :: Var.Canonical
477backwardApply =
478 Var.inCore ["Basics"] "<|"
481-- left-associative ((f >> g) >> h)
482forwardCompose :: Var.Canonical
483forwardCompose =
484 Var.inCore ["Basics"] ">>"
487-- right-associative (h << (g << f))
488backwardCompose :: Var.Canonical
489backwardCompose =
490 Var.inCore ["Basics"] "<<"
493type Assoc =
494 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
497left :: Assoc
498left =
499 flip (,)
502right :: Assoc
503right =
504 (,)
507collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
508collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
509 case expr of
510 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
511 let
512 (func, arg) = assoc leftExpr rightExpr
513 in
514 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
516 _ ->
517 makeRoot annExpr