git.iko.soy/elk
on d4e067d help sign in
elk/src/Optimize.hs
453 lines 12 kB history raw
1module Optimize (optimize) where
3import Control.Applicative
4import qualified Control.Monad as M
5import qualified Control.Monad.State as State
6import qualified Data.Traversable as T
8import qualified AST.Expression.General as Expr
9import qualified AST.Expression.Canonical as Can
10import qualified AST.Expression.Optimized as Opt
11import qualified AST.Module as Module
12import qualified AST.Module.Name as ModuleName
13import qualified AST.Pattern as P
14import qualified AST.Variable as Var
15import qualified Optimize.Environment as Env
16import qualified Optimize.Patterns as Patterns
17import qualified Reporting.Annotation as A
18import qualified Reporting.Region as R
21-- OPTIMIZE - turn a canonical module into a list of optimized defs
23optimize :: Module.CanonicalModule -> Module.Optimized
24optimize module_ =
25 let
26 home =
27 Module.name module_
29 (defs, _) =
30 flattenLets [] (Module.program (Module.body module_))
31 in
32 State.evalState
33 (concat <$> mapM (optimizeDef (Just home)) defs)
34 (Env.Env False 0)
37flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
38flattenLets defs annExpr@(A.A _ expr) =
39 case expr of
40 Expr.Let ds body ->
41 flattenLets (defs ++ ds) body
43 _ ->
44 (defs, annExpr)
47-- CONVERT DEFINITIONS
49optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Env.Optimizer [Opt.Def]
50optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
51 let
52 (args, rawBody) =
53 Expr.collectLambdas expression
54 in
55 case (ptrn, args) of
56 (P.Var name, _ : _) ->
58 do (Env.Env htc uid) <- State.get
59 State.put (Env.Env False uid)
61 (args, body) <- optimizeFunction (Just name) args rawBody
63 (Env.Env isTailCall uid') <- State.get
64 State.put (Env.Env htc uid')
66 let facts = Opt.Facts home deps
68 return $
69 if isTailCall then
71 [ Opt.TailDef facts name args body ]
73 else
75 [ Opt.Def facts name (Opt.Function args body) ]
77 (P.Var name, []) ->
79 do hasTailCall <- State.gets Env._hasTailCall
80 body <- optimizeExpr Nothing rawBody
81 Env.setTailCall hasTailCall
82 return [ Opt.Def (Opt.Facts home deps) name body ]
84 (_, []) ->
86 do name <- Env.freshName
87 optBody <- optimizeExpr Nothing rawBody
88 let optDef = Opt.Def (Opt.Facts home deps) name optBody
89 return (optDef : patternToDefs home name pattern)
91 _ ->
92 error "there should never be a function where the name is not a P.Var"
95-- OPTIMIZE FUNCTION
97optimizeFunction
98 :: Maybe String
99 -> [P.CanonicalPattern]
100 -> Can.Expr
101 -> Env.Optimizer ([String], Opt.Expr)
102optimizeFunction maybeName patterns givenBody =
103 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
104 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
105 return (argNames, optBody)
108makeContext :: [String] -> String -> Context
109makeContext argNames name =
110 Just (name, length argNames, argNames)
113depatternArgs
114 :: ([String], Can.Expr)
115 -> P.CanonicalPattern
116 -> Env.Optimizer ([String], Can.Expr)
117depatternArgs (names, rawExpr) ptrn =
118 do (name, expr) <- depattern ptrn rawExpr
119 return (name : names, expr)
122patternToDefs
123 :: Maybe ModuleName.Canonical
124 -> String
125 -> P.CanonicalPattern
126 -> [Opt.Def]
127patternToDefs home rootName pattern =
128 let
129 (deps, rootExpr) =
130 case home of
131 Nothing ->
132 ( []
133 , Opt.Var (Var.local rootName)
134 )
136 Just moduleName ->
137 ( [ Var.TopLevelVar moduleName rootName ]
138 , Opt.Var (Var.topLevel moduleName rootName)
139 )
141 toDef (name, accessExpr) =
142 Opt.Def (Opt.Facts home deps) name accessExpr
143 in
144 map toDef (substitutions rootExpr pattern)
147substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
148substitutions expr (A.A _ pattern) =
149 case pattern of
150 P.Var name ->
151 [ (name, expr) ]
153 P.Literal _ ->
154 []
156 P.Record fields ->
157 map (\name -> (name, Opt.Access expr name)) fields
159 P.Anything ->
160 []
162 P.Alias alias realPattern ->
163 (alias, expr) : substitutions expr realPattern
165 P.Data _ patterns ->
166 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
169-- OPTIMIZE EXPRESSIONS
171type Context = Maybe (String, Int, [String])
174{-| This optimizer detects tail calls. The `Context` holds a function name,
175function arity, and list of argument names. If a tail call is possible, the
176`Context` will be filled in and we check any function we run into. If we go
177down a route that cannot possibly have tail calls, we remove the context.
178-}
179optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
180optimizeExpr context annExpr@(A.A region expression) =
181 let
182 keepLooking =
183 optimizeExpr context
184 -- this *might* find a tail call, use this when something is a potential
185 -- tail call
187 justConvert =
188 optimizeExpr Nothing
189 -- this will never find tail calls, use this when you are working with
190 -- expressions that cannot be turned into tail calls
191 in
192 case expression of
193 Expr.Literal lit ->
194 pure (Opt.Literal lit)
196 Expr.Var name ->
197 if name == Var.inCore ["Debug"] "crash" then
198 pure (Opt.Crash region Nothing)
200 else
201 pure (Opt.Var name)
203 Expr.Range lowExpr highExpr ->
204 Opt.Range
205 <$> justConvert lowExpr
206 <*> justConvert highExpr
208 Expr.ExplicitList elements ->
209 Opt.ExplicitList
210 <$> T.traverse justConvert elements
212 Expr.Binop op leftExpr rightExpr ->
213 optimizeBinop context region op leftExpr rightExpr
215 Expr.Lambda _ _ ->
216 let
217 (patterns, body) =
218 Expr.collectLambdas annExpr
219 in
220 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
222 Expr.App _ _ ->
223 let
224 (func:args) =
225 Expr.collectApps annExpr
226 in
227 case isTailCall context func (length args) of
228 Just (name, argNames) ->
229 do optArgs <- T.traverse justConvert args
230 Env.setTailCall True
231 return (Opt.TailCall name argNames optArgs)
233 Nothing ->
234 do hasTailCall <- State.gets Env._hasTailCall
235 optFunc <- justConvert func
236 optArgs <- T.traverse justConvert args
237 Env.setTailCall hasTailCall
238 return (Opt.Call optFunc optArgs)
240 Expr.If branches finally ->
241 let
242 crawlBranch (cond,branch) =
243 (,) <$> justConvert cond <*> keepLooking branch
244 in
245 Opt.If
246 <$> T.traverse crawlBranch branches
247 <*> keepLooking finally
249 Expr.Let defs body ->
250 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
251 Opt.Let optDefs <$> keepLooking body
253 Expr.Case expr branches ->
254 do optExpr <- optimizeExpr Nothing expr
255 optBranches <- T.traverse (mapSnd (optimizeExpr context)) branches
256 Patterns.optimize region optExpr optBranches
258 Expr.Data name exprs ->
259 Opt.Data name <$> T.traverse justConvert exprs
261 Expr.Access record field ->
262 Opt.Access
263 <$> justConvert record
264 <*> pure field
266 Expr.Update record fields ->
267 Opt.Update
268 <$> justConvert record
269 <*> T.traverse (mapSnd justConvert) fields
271 Expr.Record fields ->
272 Opt.Record
273 <$> T.traverse (mapSnd justConvert) fields
275 Expr.GLShader uid src gltipe ->
276 pure (Opt.GLShader uid src gltipe)
278 Expr.Port impl ->
279 Opt.Port <$>
280 case impl of
281 Expr.In name tipe ->
282 pure (Expr.In name tipe)
284 Expr.Out name expr tipe ->
285 (\e -> Expr.Out name e tipe) <$> justConvert expr
287 Expr.Task name expr tipe ->
288 (\e -> Expr.Task name e tipe) <$> justConvert expr
291mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
292mapSnd func (x, a) =
293 (,) x <$> func a
296-- DETECT TAIL CALL
298isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
299isTailCall context (A.A _ function) arity =
300 case (context, function) of
301 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
303 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
304 Just (name, argNames)
306 else
307 Nothing
309 _ ->
310 Nothing
313-- DEPATTERN
314-- given a pattern and an expression, push the actual pattern matching into the
315-- expression as a case-expression if necessary. The pattern compiler will work
316-- it out from there
318depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
319depattern canPattern@(A.A _ pattern) expr =
320 let
321 ann =
322 A.A (error "the annotation added in 'depattern' should not be observed!")
324 caseExpr e branches =
325 ann (Expr.Case e branches)
327 var name =
328 ann (Expr.Var (Var.local name))
330 caseify =
331 do name <- Env.freshName
332 return
333 ( name
334 , caseExpr (var name) [ (canPattern, expr) ]
335 )
336 in
337 case pattern of
338 P.Var name ->
339 return (name, expr)
341 P.Literal _ ->
342 caseify
344 P.Record _ ->
345 caseify
347 P.Anything ->
348 do name <- Env.freshName
349 return (name, expr)
351 P.Alias _ _ ->
352 caseify
354 P.Data _ _ ->
355 caseify
358-- OPTIMIZE BINOPS
360optimizeBinop
361 :: Context
362 -> R.Region
363 -> Var.Canonical
364 -> Can.Expr
365 -> Can.Expr
366 -> Env.Optimizer Opt.Expr
367optimizeBinop context region op leftExpr rightExpr =
368 let
369 ann = A.A region
370 binop = ann (Expr.Binop op leftExpr rightExpr)
371 in
372 if op == forwardApply then
374 optimizeExpr context (collect id left forwardApply binop)
376 else if op == backwardApply then
378 optimizeExpr context (collect id right backwardApply binop)
380 else if op == forwardCompose then
382 do var <- Env.freshName
383 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
384 let body = collect makeRoot left forwardCompose binop
385 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
387 else if op == backwardCompose then
389 do var <- Env.freshName
390 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
391 let body = collect makeRoot right backwardCompose binop
392 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
394 else if op == Var.Canonical Var.BuiltIn "::" then
396 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
398 else
400 Opt.Binop op
401 <$> optimizeExpr Nothing leftExpr
402 <*> optimizeExpr Nothing rightExpr
405-- left-associative ((x |> f) |> g)
406forwardApply :: Var.Canonical
407forwardApply =
408 Var.inCore ["Basics"] "|>"
411-- right-associative (g <| (f <| x))
412backwardApply :: Var.Canonical
413backwardApply =
414 Var.inCore ["Basics"] "<|"
417-- left-associative ((f >> g) >> h)
418forwardCompose :: Var.Canonical
419forwardCompose =
420 Var.inCore ["Basics"] ">>"
423-- right-associative (h << (g << f))
424backwardCompose :: Var.Canonical
425backwardCompose =
426 Var.inCore ["Basics"] "<<"
429type Assoc =
430 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
433left :: Assoc
434left =
435 flip (,)
438right :: Assoc
439right =
440 (,)
443collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
444collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
445 case expr of
446 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
447 let
448 (func, arg) = assoc leftExpr rightExpr
449 in
450 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
452 _ ->
453 makeRoot annExpr