git.iko.soy/elk
on bc71d83 help sign in
elk/src/Optimize.hs
477 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 as Module
11import qualified AST.Module.Name as ModuleName
12import qualified AST.Pattern as P
13import qualified AST.Variable as Var
14import qualified Optimize.Environment as Env
15import qualified Optimize.Inline as Inline
16import qualified Optimize.Patterns as Patterns
17import qualified Optimize.Patterns.DecisionTree as DT
18import qualified Reporting.Annotation as A
19import qualified Reporting.Region as R
22-- OPTIMIZE
24optimize :: DT.VariantDict -> Module.CanonicalModule -> Module.Optimized
25optimize variantDict module_@(Module.Module home _ _ _ _ body) =
26 let
27 (defs, _) =
28 flattenLets [] (Module.program body)
30 optDefs =
31 Env.run variantDict (concat <$> mapM (optimizeDef (Just home)) defs)
32 in
33 module_ { Module.body = body { Module.program = optDefs } }
36flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
37flattenLets defs annExpr@(A.A _ expr) =
38 case expr of
39 Expr.Let ds body ->
40 flattenLets (defs ++ ds) body
42 _ ->
43 (defs, annExpr)
46-- CONVERT DEFINITIONS
48optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Env.Optimizer [Opt.Def]
49optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
50 let
51 (args, rawBody) =
52 Expr.collectLambdas expression
53 in
54 case (ptrn, args) of
55 (P.Var name, _ : _) ->
57 do htc <- Env.getTailCall
58 Env.setTailCall False
60 (args, body) <- optimizeFunction (Just name) args rawBody
62 isTailCall <- Env.getTailCall
63 Env.setTailCall htc
65 let facts = Opt.Facts home deps
67 return $
68 if isTailCall then
70 [ Opt.TailDef facts name args body ]
72 else
74 [ Opt.Def facts name (Opt.Function args body) ]
76 (P.Var name, []) ->
78 do hasTailCall <- Env.getTailCall
79 body <- optimizeExpr Nothing rawBody
80 Env.setTailCall hasTailCall
81 return [ Opt.Def (Opt.Facts home deps) name body ]
83 (_, []) ->
85 do name <- Env.freshName
86 optBody <- optimizeExpr Nothing rawBody
87 let optDef = Opt.Def (Opt.Facts home deps) name optBody
88 return (optDef : patternToDefs home name pattern)
90 _ ->
91 error "there should never be a function where the name is not a P.Var"
94-- OPTIMIZE FUNCTION
96optimizeFunction
97 :: Maybe String
98 -> [P.CanonicalPattern]
99 -> Can.Expr
100 -> Env.Optimizer ([String], Opt.Expr)
101optimizeFunction maybeName patterns givenBody =
102 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
103 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
104 return (argNames, optBody)
107makeContext :: [String] -> String -> Context
108makeContext argNames name =
109 Just (name, length argNames, argNames)
112depatternArgs
113 :: ([String], Can.Expr)
114 -> P.CanonicalPattern
115 -> Env.Optimizer ([String], Can.Expr)
116depatternArgs (names, rawExpr) ptrn =
117 do (name, expr) <- depattern ptrn rawExpr
118 return (name : names, expr)
121patternToDefs
122 :: Maybe ModuleName.Canonical
123 -> String
124 -> P.CanonicalPattern
125 -> [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)
146-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
148patternToSubstitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
149patternToSubstitutions expr (A.A _ pattern) =
150 case pattern of
151 P.Var name ->
152 [ (name, expr) ]
154 P.Literal _ ->
155 []
157 P.Record fields ->
158 map (\name -> (name, Opt.Access expr name)) fields
160 P.Anything ->
161 []
163 P.Alias alias realPattern ->
164 (alias, expr) : patternToSubstitutions expr realPattern
166 P.Data _ patterns ->
167 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
170-- OPTIMIZE EXPRESSIONS
172type Context = Maybe (String, Int, [String])
175{-| This optimizer detects tail calls. The `Context` holds a function name,
176function arity, and list of argument names. If a tail call is possible, the
177`Context` will be filled in and we check any function we run into. If we go
178down a route that cannot possibly have tail calls, we remove the context.
179-}
180optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
181optimizeExpr context annExpr@(A.A region expression) =
182 let
183 keepLooking =
184 optimizeExpr context
185 -- this *might* find a tail call, use this when something is a potential
186 -- tail call
188 justConvert =
189 optimizeExpr Nothing
190 -- this will never find tail calls, use this when you are working with
191 -- expressions that cannot be turned into tail calls
192 in
193 case expression of
194 Expr.Literal lit ->
195 pure (Opt.Literal lit)
197 Expr.Var name ->
198 if name == Var.inCore ["Debug"] "crash" then
199 pure (Opt.Crash region Nothing)
201 else
202 pure (Opt.Var name)
204 Expr.Range lowExpr highExpr ->
205 Opt.Range
206 <$> justConvert lowExpr
207 <*> justConvert highExpr
209 Expr.ExplicitList elements ->
210 Opt.ExplicitList
211 <$> T.traverse justConvert elements
213 Expr.Binop op leftExpr rightExpr ->
214 optimizeBinop context region op leftExpr rightExpr
216 Expr.Lambda _ _ ->
217 let
218 (patterns, body) =
219 Expr.collectLambdas annExpr
220 in
221 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
223 Expr.App _ _ ->
224 let
225 (func:args) =
226 Expr.collectApps annExpr
227 in
228 case isTailCall context func (length args) of
229 Just (name, argNames) ->
230 do optArgs <- T.traverse justConvert args
231 Env.setTailCall True
232 return (Opt.TailCall name argNames optArgs)
234 Nothing ->
235 do hasTailCall <- Env.getTailCall
236 optFunc <- justConvert func
237 optArgs <- T.traverse justConvert args
238 Env.setTailCall hasTailCall
239 return (Opt.Call optFunc optArgs)
241 Expr.If branches finally ->
242 let
243 crawlBranch (cond,branch) =
244 (,) <$> justConvert cond <*> keepLooking branch
245 in
246 Opt.If
247 <$> T.traverse crawlBranch branches
248 <*> keepLooking finally
250 Expr.Let defs body ->
251 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
252 Opt.Let optDefs <$> keepLooking body
254 Expr.Case expr branches ->
255 do optExpr <- optimizeExpr Nothing expr
256 variantDict <- Env.getVariantDict
258 name <- Env.freshName
259 optBranches <- T.traverse (optimizeBranch context name) branches
260 optCase <- Patterns.optimize variantDict region name optBranches
261 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
263 Expr.Data name exprs ->
264 Opt.Data name <$> T.traverse justConvert exprs
266 Expr.Access record field ->
267 Opt.Access
268 <$> justConvert record
269 <*> pure field
271 Expr.Update record fields ->
272 Opt.Update
273 <$> justConvert record
274 <*> T.traverse (mapSnd justConvert) fields
276 Expr.Record fields ->
277 Opt.Record
278 <$> T.traverse (mapSnd justConvert) fields
280 Expr.GLShader uid src gltipe ->
281 pure (Opt.GLShader uid src gltipe)
283 Expr.Port impl ->
284 Opt.Port <$>
285 case impl of
286 Expr.In name tipe ->
287 pure (Expr.In name tipe)
289 Expr.Out name expr tipe ->
290 (\e -> Expr.Out name e tipe) <$> justConvert expr
292 Expr.Task name expr tipe ->
293 (\e -> Expr.Task name e tipe) <$> justConvert expr
296mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
297mapSnd func (x, a) =
298 (,) x <$> func a
301-- OPTIMIZE CASE BRANCHES
303optimizeBranch
304 :: Context
305 -> String
306 -> (P.CanonicalPattern, Can.Expr)
307 -> Env.Optimizer (P.CanonicalPattern, Opt.Expr)
308optimizeBranch context exprName (pattern, canonicalExpr) =
309 let
310 root =
311 Opt.Var (Var.Canonical Var.Local exprName)
313 substitutions =
314 Map.fromList (patternToSubstitutions root pattern)
315 in
316 do optExpr <- optimizeExpr context canonicalExpr
317 (,) pattern <$> Inline.inline substitutions optExpr
320-- DETECT TAIL CALL
322isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
323isTailCall context (A.A _ function) arity =
324 case (context, function) of
325 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
327 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
328 Just (name, argNames)
330 else
331 Nothing
333 _ ->
334 Nothing
337-- DEPATTERN
338-- given a pattern and an expression, push the actual pattern matching into the
339-- expression as a case-expression if necessary. The pattern compiler will work
340-- it out from there
342depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
343depattern canPattern@(A.A _ pattern) expr =
344 let
345 ann =
346 A.A (error "the annotation added in 'depattern' should not be observed!")
348 caseExpr e branches =
349 ann (Expr.Case e branches)
351 var name =
352 ann (Expr.Var (Var.local name))
354 caseify =
355 do name <- Env.freshName
356 return
357 ( name
358 , caseExpr (var name) [ (canPattern, expr) ]
359 )
360 in
361 case pattern of
362 P.Var name ->
363 return (name, expr)
365 P.Literal _ ->
366 caseify
368 P.Record _ ->
369 caseify
371 P.Anything ->
372 do name <- Env.freshName
373 return (name, expr)
375 P.Alias _ _ ->
376 caseify
378 P.Data _ _ ->
379 caseify
382-- OPTIMIZE BINOPS
384optimizeBinop
385 :: Context
386 -> R.Region
387 -> Var.Canonical
388 -> Can.Expr
389 -> Can.Expr
390 -> Env.Optimizer Opt.Expr
391optimizeBinop context region op leftExpr rightExpr =
392 let
393 ann = A.A region
394 binop = ann (Expr.Binop op leftExpr rightExpr)
395 in
396 if op == forwardApply then
398 optimizeExpr context (collect id left forwardApply binop)
400 else if op == backwardApply then
402 optimizeExpr context (collect id right backwardApply binop)
404 else if op == forwardCompose then
406 do var <- Env.freshName
407 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
408 let body = collect makeRoot left forwardCompose binop
409 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
411 else if op == backwardCompose then
413 do var <- Env.freshName
414 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
415 let body = collect makeRoot right backwardCompose binop
416 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
418 else if op == Var.Canonical Var.BuiltIn "::" then
420 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
422 else
424 Opt.Binop op
425 <$> optimizeExpr Nothing leftExpr
426 <*> optimizeExpr Nothing rightExpr
429-- left-associative ((x |> f) |> g)
430forwardApply :: Var.Canonical
431forwardApply =
432 Var.inCore ["Basics"] "|>"
435-- right-associative (g <| (f <| x))
436backwardApply :: Var.Canonical
437backwardApply =
438 Var.inCore ["Basics"] "<|"
441-- left-associative ((f >> g) >> h)
442forwardCompose :: Var.Canonical
443forwardCompose =
444 Var.inCore ["Basics"] ">>"
447-- right-associative (h << (g << f))
448backwardCompose :: Var.Canonical
449backwardCompose =
450 Var.inCore ["Basics"] "<<"
453type Assoc =
454 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
457left :: Assoc
458left =
459 flip (,)
462right :: Assoc
463right =
464 (,)
467collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
468collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
469 case expr of
470 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
471 let
472 (func, arg) = assoc leftExpr rightExpr
473 in
474 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
476 _ ->
477 makeRoot annExpr