git.iko.soy/elk
on 563d322 help sign in
elk/src/Optimize.hs
452 lines 12 kB history raw
1module Optimize (optimize) where
3import Control.Applicative
4import qualified Control.Monad as M
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.Patterns as Patterns
16import qualified Optimize.Patterns.DecisionTree as DT
17import qualified Reporting.Annotation as A
18import qualified Reporting.Region as R
21-- OPTIMIZE
23optimize :: DT.VariantDict -> Module.CanonicalModule -> Module.Optimized
24optimize variantDict module_@(Module.Module home _ _ _ _ body) =
25 let
26 (defs, _) =
27 flattenLets [] (Module.program body)
29 optDefs =
30 Env.run variantDict (concat <$> mapM (optimizeDef (Just home)) defs)
31 in
32 module_ { Module.body = body { Module.program = optDefs } }
35flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
36flattenLets defs annExpr@(A.A _ expr) =
37 case expr of
38 Expr.Let ds body ->
39 flattenLets (defs ++ ds) body
41 _ ->
42 (defs, annExpr)
45-- CONVERT DEFINITIONS
47optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Env.Optimizer [Opt.Def]
48optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
49 let
50 (args, rawBody) =
51 Expr.collectLambdas expression
52 in
53 case (ptrn, args) of
54 (P.Var name, _ : _) ->
56 do htc <- Env.getTailCall
57 Env.setTailCall False
59 (args, body) <- optimizeFunction (Just name) args rawBody
61 isTailCall <- Env.getTailCall
62 Env.setTailCall htc
64 let facts = Opt.Facts home deps
66 return $
67 if isTailCall then
69 [ Opt.TailDef facts name args body ]
71 else
73 [ Opt.Def facts name (Opt.Function args body) ]
75 (P.Var name, []) ->
77 do hasTailCall <- Env.getTailCall
78 body <- optimizeExpr Nothing rawBody
79 Env.setTailCall hasTailCall
80 return [ Opt.Def (Opt.Facts home deps) name body ]
82 (_, []) ->
84 do name <- Env.freshName
85 optBody <- optimizeExpr Nothing rawBody
86 let optDef = Opt.Def (Opt.Facts home deps) name optBody
87 return (optDef : patternToDefs home name pattern)
89 _ ->
90 error "there should never be a function where the name is not a P.Var"
93-- OPTIMIZE FUNCTION
95optimizeFunction
96 :: Maybe String
97 -> [P.CanonicalPattern]
98 -> Can.Expr
99 -> Env.Optimizer ([String], Opt.Expr)
100optimizeFunction maybeName patterns givenBody =
101 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
102 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
103 return (argNames, optBody)
106makeContext :: [String] -> String -> Context
107makeContext argNames name =
108 Just (name, length argNames, argNames)
111depatternArgs
112 :: ([String], Can.Expr)
113 -> P.CanonicalPattern
114 -> Env.Optimizer ([String], Can.Expr)
115depatternArgs (names, rawExpr) ptrn =
116 do (name, expr) <- depattern ptrn rawExpr
117 return (name : names, expr)
120patternToDefs
121 :: Maybe ModuleName.Canonical
122 -> String
123 -> P.CanonicalPattern
124 -> [Opt.Def]
125patternToDefs home rootName pattern =
126 let
127 (deps, rootExpr) =
128 case home of
129 Nothing ->
130 ( []
131 , Opt.Var (Var.local rootName)
132 )
134 Just moduleName ->
135 ( [ Var.TopLevelVar moduleName rootName ]
136 , Opt.Var (Var.topLevel moduleName rootName)
137 )
139 toDef (name, accessExpr) =
140 Opt.Def (Opt.Facts home deps) name accessExpr
141 in
142 map toDef (substitutions rootExpr pattern)
145substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
146substitutions expr (A.A _ pattern) =
147 case pattern of
148 P.Var name ->
149 [ (name, expr) ]
151 P.Literal _ ->
152 []
154 P.Record fields ->
155 map (\name -> (name, Opt.Access expr name)) fields
157 P.Anything ->
158 []
160 P.Alias alias realPattern ->
161 (alias, expr) : substitutions expr realPattern
163 P.Data _ patterns ->
164 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
167-- OPTIMIZE EXPRESSIONS
169type Context = Maybe (String, Int, [String])
172{-| This optimizer detects tail calls. The `Context` holds a function name,
173function arity, and list of argument names. If a tail call is possible, the
174`Context` will be filled in and we check any function we run into. If we go
175down a route that cannot possibly have tail calls, we remove the context.
176-}
177optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
178optimizeExpr context annExpr@(A.A region expression) =
179 let
180 keepLooking =
181 optimizeExpr context
182 -- this *might* find a tail call, use this when something is a potential
183 -- tail call
185 justConvert =
186 optimizeExpr Nothing
187 -- this will never find tail calls, use this when you are working with
188 -- expressions that cannot be turned into tail calls
189 in
190 case expression of
191 Expr.Literal lit ->
192 pure (Opt.Literal lit)
194 Expr.Var name ->
195 if name == Var.inCore ["Debug"] "crash" then
196 pure (Opt.Crash region Nothing)
198 else
199 pure (Opt.Var name)
201 Expr.Range lowExpr highExpr ->
202 Opt.Range
203 <$> justConvert lowExpr
204 <*> justConvert highExpr
206 Expr.ExplicitList elements ->
207 Opt.ExplicitList
208 <$> T.traverse justConvert elements
210 Expr.Binop op leftExpr rightExpr ->
211 optimizeBinop context region op leftExpr rightExpr
213 Expr.Lambda _ _ ->
214 let
215 (patterns, body) =
216 Expr.collectLambdas annExpr
217 in
218 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
220 Expr.App _ _ ->
221 let
222 (func:args) =
223 Expr.collectApps annExpr
224 in
225 case isTailCall context func (length args) of
226 Just (name, argNames) ->
227 do optArgs <- T.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 <- T.traverse justConvert args
235 Env.setTailCall hasTailCall
236 return (Opt.Call optFunc optArgs)
238 Expr.If branches finally ->
239 let
240 crawlBranch (cond,branch) =
241 (,) <$> justConvert cond <*> keepLooking branch
242 in
243 Opt.If
244 <$> T.traverse crawlBranch branches
245 <*> keepLooking finally
247 Expr.Let defs body ->
248 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
249 Opt.Let optDefs <$> keepLooking body
251 Expr.Case expr branches ->
252 do optExpr <- optimizeExpr Nothing expr
253 optBranches <- T.traverse (mapSnd (optimizeExpr context)) branches
254 variantDict <- Env.getVariantDict
255 Patterns.optimize variantDict region optExpr optBranches
257 Expr.Data name exprs ->
258 Opt.Data name <$> T.traverse justConvert exprs
260 Expr.Access record field ->
261 Opt.Access
262 <$> justConvert record
263 <*> pure field
265 Expr.Update record fields ->
266 Opt.Update
267 <$> justConvert record
268 <*> T.traverse (mapSnd justConvert) fields
270 Expr.Record fields ->
271 Opt.Record
272 <$> T.traverse (mapSnd justConvert) fields
274 Expr.GLShader uid src gltipe ->
275 pure (Opt.GLShader uid src gltipe)
277 Expr.Port impl ->
278 Opt.Port <$>
279 case impl of
280 Expr.In name tipe ->
281 pure (Expr.In name tipe)
283 Expr.Out name expr tipe ->
284 (\e -> Expr.Out name e tipe) <$> justConvert expr
286 Expr.Task name expr tipe ->
287 (\e -> Expr.Task name e tipe) <$> justConvert expr
290mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
291mapSnd func (x, a) =
292 (,) x <$> func a
295-- DETECT TAIL CALL
297isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
298isTailCall context (A.A _ function) arity =
299 case (context, function) of
300 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
302 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
303 Just (name, argNames)
305 else
306 Nothing
308 _ ->
309 Nothing
312-- DEPATTERN
313-- given a pattern and an expression, push the actual pattern matching into the
314-- expression as a case-expression if necessary. The pattern compiler will work
315-- it out from there
317depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
318depattern canPattern@(A.A _ pattern) expr =
319 let
320 ann =
321 A.A (error "the annotation added in 'depattern' should not be observed!")
323 caseExpr e branches =
324 ann (Expr.Case e branches)
326 var name =
327 ann (Expr.Var (Var.local name))
329 caseify =
330 do name <- Env.freshName
331 return
332 ( name
333 , caseExpr (var name) [ (canPattern, expr) ]
334 )
335 in
336 case pattern of
337 P.Var name ->
338 return (name, expr)
340 P.Literal _ ->
341 caseify
343 P.Record _ ->
344 caseify
346 P.Anything ->
347 do name <- Env.freshName
348 return (name, expr)
350 P.Alias _ _ ->
351 caseify
353 P.Data _ _ ->
354 caseify
357-- OPTIMIZE BINOPS
359optimizeBinop
360 :: Context
361 -> R.Region
362 -> Var.Canonical
363 -> Can.Expr
364 -> Can.Expr
365 -> Env.Optimizer Opt.Expr
366optimizeBinop context region op leftExpr rightExpr =
367 let
368 ann = A.A region
369 binop = ann (Expr.Binop op leftExpr rightExpr)
370 in
371 if op == forwardApply then
373 optimizeExpr context (collect id left forwardApply binop)
375 else if op == backwardApply then
377 optimizeExpr context (collect id right backwardApply binop)
379 else if op == forwardCompose then
381 do var <- Env.freshName
382 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
383 let body = collect makeRoot left forwardCompose binop
384 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
386 else if op == backwardCompose then
388 do var <- Env.freshName
389 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
390 let body = collect makeRoot right backwardCompose binop
391 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
393 else if op == Var.Canonical Var.BuiltIn "::" then
395 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
397 else
399 Opt.Binop op
400 <$> optimizeExpr Nothing leftExpr
401 <*> optimizeExpr Nothing rightExpr
404-- left-associative ((x |> f) |> g)
405forwardApply :: Var.Canonical
406forwardApply =
407 Var.inCore ["Basics"] "|>"
410-- right-associative (g <| (f <| x))
411backwardApply :: Var.Canonical
412backwardApply =
413 Var.inCore ["Basics"] "<|"
416-- left-associative ((f >> g) >> h)
417forwardCompose :: Var.Canonical
418forwardCompose =
419 Var.inCore ["Basics"] ">>"
422-- right-associative (h << (g << f))
423backwardCompose :: Var.Canonical
424backwardCompose =
425 Var.inCore ["Basics"] "<<"
428type Assoc =
429 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
432left :: Assoc
433left =
434 flip (,)
437right :: Assoc
438right =
439 (,)
442collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
443collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
444 case expr of
445 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
446 let
447 (func, arg) = assoc leftExpr rightExpr
448 in
449 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
451 _ ->
452 makeRoot annExpr