@@ -1,9 +1,487 @@
1−module Optimize where
1+module Optimize (optimize) where
22
3−import qualified AST.Module as M
4−import qualified Optimize.TailCalls as TC
3+import Control.Applicative
4+import qualified Control.Monad as M
5+import qualified Control.Monad.State as State
6+import qualified Data.Map as Map
7+import qualified Data.Traversable as T
58
9+import qualified AST.Expression.General as Expr
10+import qualified AST.Expression.Canonical as Can
11+import qualified AST.Expression.Optimized as Opt
12+import qualified AST.Module as Module
13+import qualified AST.Module.Name as ModuleName
14+import qualified AST.Pattern as P
15+import qualified AST.Variable as Var
16+import qualified Optimize.Cases as Cases
17+import qualified Reporting.Annotation as A
18+import qualified Reporting.Region as R
619
7−optimize :: M.CanonicalModule -> M.Optimized
8−optimize =
9− TC.optimize
20+
21+-- OPTIMIZE FOR TAIL CALLS
22+
23+optimize :: Module.CanonicalModule -> Module.Optimized
24+optimize module_ =
25+ let
26+ home =
27+ Module.name module_
28+
29+ (defs, _) =
30+ flattenLets [] (Module.program (Module.body module_))
31+ in
32+ State.evalState
33+ (concat <$> mapM (definition (Just home)) defs)
34+ (Env False 0)
35+
36+
37+flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
38+flattenLets defs annExpr@(A.A _ expr) =
39+ case expr of
40+ Expr.Let ds body ->
41+ flattenLets (defs ++ ds) body
42+
43+ _ ->
44+ (defs, annExpr)
45+
46+
47+-- TRACKER TO THREAD FINDINGS THROUGH TRAVERSAL
48+
49+type Optimizer a =
50+ State.State Env a
51+
52+
53+data Env = Env
54+ { _hasTailCall :: Bool
55+ , _uid :: Int
56+ }
57+
58+
59+setTailCall :: Bool -> Optimizer ()
60+setTailCall bool =
61+ do uid <- State.gets _uid
62+ State.put (Env bool uid)
63+
64+
65+mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
66+mapSnd func (x, a) =
67+ (,) x <$> func a
68+
69+
70+-- CONVERT DEFINITIONS
71+
72+definition :: Maybe ModuleName.Canonical -> Can.Def -> Optimizer [Opt.Def]
73+definition home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
74+ let
75+ (args, rawBody) =
76+ Expr.collectLambdas expression
77+ in
78+ case (ptrn, args) of
79+ (P.Var name, _ : _) ->
80+
81+ do (Env htc uid) <- State.get
82+ State.put (Env False uid)
83+
84+ (args, body) <- optimizeFunction (Just name) args rawBody
85+
86+ (Env isTailCall uid') <- State.get
87+ State.put (Env htc uid')
88+
89+ let facts = Opt.Facts home deps
90+
91+ return $
92+ if isTailCall then
93+
94+ [ Opt.TailDef facts name args body ]
95+
96+ else
97+
98+ [ Opt.Def facts name (Opt.Function args body) ]
99+
100+ (P.Var name, []) ->
101+
102+ do hasTailCall <- State.gets _hasTailCall
103+ body <- findTailCalls Nothing rawBody
104+ setTailCall hasTailCall
105+ return [ Opt.Def (Opt.Facts home deps) name body ]
106+
107+ (_, []) ->
108+
109+ do name <- freshName
110+ optBody <- findTailCalls Nothing rawBody
111+ let optDef = Opt.Def (Opt.Facts home deps) name optBody
112+ return (optDef : patternToDefs home name pattern)
113+
114+ _ ->
115+ error "there should never be a function where the name is not a P.Var"
116+
117+
118+-- OPTIMIZE FUNCTION
119+
120+optimizeFunction
121+ :: Maybe String
122+ -> [P.CanonicalPattern]
123+ -> Can.Expr
124+ -> Optimizer ([String], Opt.Expr)
125+optimizeFunction maybeName patterns givenBody =
126+ do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
127+ optBody <- findTailCalls (makeContext argNames =<< maybeName) body
128+ return (argNames, optBody)
129+
130+
131+makeContext :: [String] -> String -> Context
132+makeContext argNames name =
133+ Just (name, length argNames, argNames)
134+
135+
136+depatternArgs
137+ :: ([String], Can.Expr)
138+ -> P.CanonicalPattern
139+ -> Optimizer ([String], Can.Expr)
140+depatternArgs (names, rawExpr) ptrn =
141+ do (name, expr) <- depattern ptrn rawExpr
142+ return (name : names, expr)
143+
144+
145+patternToDefs
146+ :: Maybe ModuleName.Canonical
147+ -> String
148+ -> P.CanonicalPattern
149+ -> [Opt.Def]
150+patternToDefs home rootName pattern =
151+ let
152+ (deps, rootExpr) =
153+ case home of
154+ Nothing ->
155+ ( []
156+ , Opt.Var (Var.local rootName)
157+ )
158+
159+ Just moduleName ->
160+ ( [ Var.TopLevelVar moduleName rootName ]
161+ , Opt.Var (Var.topLevel moduleName rootName)
162+ )
163+
164+ toDef (name, accessExpr) =
165+ Opt.Def (Opt.Facts home deps) name accessExpr
166+ in
167+ map toDef (substitutions rootExpr pattern)
168+
169+
170+substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
171+substitutions expr (A.A _ pattern) =
172+ case pattern of
173+ P.Var name ->
174+ [ (name, expr) ]
175+
176+ P.Literal _ ->
177+ []
178+
179+ P.Record fields ->
180+ map (\name -> (name, Opt.Access expr name)) fields
181+
182+ P.Anything ->
183+ []
184+
185+ P.Alias alias realPattern ->
186+ (alias, expr) : substitutions expr realPattern
187+
188+ P.Data _ patterns ->
189+ concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
190+
191+
192+-- CONVERT EXPRESSIONS
193+
194+type Context = Maybe (String, Int, [String])
195+
196+
197+findTailCalls :: Context -> Can.Expr -> Optimizer Opt.Expr
198+findTailCalls context annExpr@(A.A region expression) =
199+ let
200+ keepLooking =
201+ findTailCalls context
202+ -- this *might* find a tail call, use this when something is a potential
203+ -- tail call
204+
205+ justConvert =
206+ findTailCalls Nothing
207+ -- this will never find tail calls, use this when you are working with
208+ -- expressions that cannot be turned into tail calls
209+ in
210+ case expression of
211+ Expr.Literal lit ->
212+ pure (Opt.Literal lit)
213+
214+ Expr.Var name ->
215+ pure (Opt.Var name)
216+
217+ Expr.Range lowExpr highExpr ->
218+ Opt.Range
219+ <$> justConvert lowExpr
220+ <*> justConvert highExpr
221+
222+ Expr.ExplicitList elements ->
223+ Opt.ExplicitList
224+ <$> T.traverse justConvert elements
225+
226+ Expr.Binop op leftExpr rightExpr ->
227+ optimizeBinop context region op leftExpr rightExpr
228+
229+ Expr.Lambda _ _ ->
230+ let
231+ (patterns, body) =
232+ Expr.collectLambdas annExpr
233+ in
234+ uncurry Opt.Function <$> optimizeFunction Nothing patterns body
235+
236+ Expr.App _ _ ->
237+ let
238+ (func:args) =
239+ Expr.collectApps annExpr
240+ in
241+ case isTailCall context func (length args) of
242+ Just (name, argNames) ->
243+ do optArgs <- T.traverse justConvert args
244+ setTailCall True
245+ return (Opt.TailCall name argNames optArgs)
246+
247+ Nothing ->
248+ do hasTailCall <- State.gets _hasTailCall
249+ optFunc <- justConvert func
250+ optArgs <- T.traverse justConvert args
251+ setTailCall hasTailCall
252+ return (Opt.Call optFunc optArgs)
253+
254+
255+ Expr.If branches finally ->
256+ let
257+ crawlBranch (cond,branch) =
258+ (,) <$> justConvert cond <*> keepLooking branch
259+ in
260+ Opt.If
261+ <$> T.traverse crawlBranch branches
262+ <*> keepLooking finally
263+
264+ Expr.Let defs body ->
265+ do optDefs <- concat <$> T.traverse (definition Nothing) defs
266+ Opt.Let optDefs <$> keepLooking body
267+
268+ Expr.Case expr rawBranches ->
269+ let
270+ indexify index (pattern, branch) =
271+ ( (pattern, index)
272+ , (index, branch)
273+ )
274+
275+ (patterns, indexedBranches) =
276+ unzip (zipWith indexify [0..] rawBranches)
277+
278+ decisionTree =
279+ Cases.compile (error "need a VariantDict") patterns
280+
281+ branches =
282+ Map.fromList indexedBranches
283+ in
284+ Opt.Case
285+ <$> justConvert expr
286+ <*> pure decisionTree
287+ <*> T.traverse keepLooking branches
288+
289+ Expr.Data name exprs ->
290+ Opt.Data name <$> T.traverse justConvert exprs
291+
292+ Expr.Access record field ->
293+ Opt.Access
294+ <$> justConvert record
295+ <*> pure field
296+
297+ Expr.Update record fields ->
298+ Opt.Update
299+ <$> justConvert record
300+ <*> T.traverse (mapSnd justConvert) fields
301+
302+ Expr.Record fields ->
303+ Opt.Record
304+ <$> T.traverse (mapSnd justConvert) fields
305+
306+ Expr.GLShader uid src gltipe ->
307+ pure (Opt.GLShader uid src gltipe)
308+
309+ Expr.Port impl ->
310+ Opt.Port <$>
311+ case impl of
312+ Expr.In name tipe ->
313+ pure (Expr.In name tipe)
314+
315+ Expr.Out name expr tipe ->
316+ (\e -> Expr.Out name e tipe) <$> justConvert expr
317+
318+ Expr.Task name expr tipe ->
319+ (\e -> Expr.Task name e tipe) <$> justConvert expr
320+
321+
322+-- DETECT TAIL CALL
323+
324+isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
325+isTailCall context (A.A _ function) arity =
326+ case (context, function) of
327+ ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
328+
329+ if name == ctxName && arity == ctxArity && Var.isLocalHome home then
330+ Just (name, argNames)
331+
332+ else
333+ Nothing
334+
335+ _ ->
336+ Nothing
337+
338+
339+
340+-- DEPATTERN
341+-- given a pattern and an expression, push the actual pattern matching into the
342+-- expression as a case-expression if necessary. The pattern compiler will work
343+-- it out from there
344+
345+depattern :: P.CanonicalPattern -> Can.Expr -> Optimizer (String, Can.Expr)
346+depattern canPattern@(A.A _ pattern) expr =
347+ let
348+ ann =
349+ A.A (error "the annotation added in 'depattern' should not be observed!")
350+
351+ caseExpr e branches =
352+ ann (Expr.Case e branches)
353+
354+ var name =
355+ ann (Expr.Var (Var.local name))
356+
357+ caseify =
358+ do name <- freshName
359+ return
360+ ( name
361+ , caseExpr (var name) [ (canPattern, expr) ]
362+ )
363+ in
364+ case pattern of
365+ P.Var name ->
366+ return (name, expr)
367+
368+ P.Literal _ ->
369+ caseify
370+
371+ P.Record _ ->
372+ caseify
373+
374+ P.Anything ->
375+ do name <- freshName
376+ return (name, expr)
377+
378+ P.Alias _ _ ->
379+ caseify
380+
381+ P.Data _ _ ->
382+ caseify
383+
384+
385+freshName :: Optimizer String
386+freshName =
387+ do (Env htc uid) <- State.get
388+ State.put (Env htc (uid + 1))
389+ return ("_p" ++ show uid)
390+
391+
392+-- OPTIMIZE BINOPS
393+
394+optimizeBinop
395+ :: Context
396+ -> R.Region
397+ -> Var.Canonical
398+ -> Can.Expr
399+ -> Can.Expr
400+ -> Optimizer Opt.Expr
401+optimizeBinop context region op leftExpr rightExpr =
402+ let
403+ ann = A.A region
404+ binop = ann (Expr.Binop op leftExpr rightExpr)
405+ in
406+ if op == forwardApply then
407+
408+ findTailCalls context (collect id left forwardApply binop)
409+
410+ else if op == backwardApply then
411+
412+ findTailCalls context (collect id right backwardApply binop)
413+
414+ else if op == forwardCompose then
415+
416+ do var <- freshName
417+ let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
418+ let body = collect makeRoot left forwardCompose binop
419+ findTailCalls context (ann (Expr.Lambda (ann (P.Var var)) body))
420+
421+ else if op == backwardCompose then
422+
423+ do var <- freshName
424+ let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
425+ let body = collect makeRoot right backwardCompose binop
426+ findTailCalls context (ann (Expr.Lambda (ann (P.Var var)) body))
427+
428+ else if op == Var.Canonical Var.BuiltIn "::" then
429+
430+ findTailCalls context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
431+
432+ else
433+
434+ Opt.Binop op
435+ <$> findTailCalls Nothing leftExpr
436+ <*> findTailCalls Nothing rightExpr
437+
438+
439+-- left-associative ((x |> f) |> g)
440+forwardApply :: Var.Canonical
441+forwardApply =
442+ Var.inCore ["Basics"] "|>"
443+
444+
445+-- right-associative (g <| (f <| x))
446+backwardApply :: Var.Canonical
447+backwardApply =
448+ Var.inCore ["Basics"] "<|"
449+
450+
451+-- left-associative ((f >> g) >> h)
452+forwardCompose :: Var.Canonical
453+forwardCompose =
454+ Var.inCore ["Basics"] ">>"
455+
456+
457+-- right-associative (h << (g << f))
458+backwardCompose :: Var.Canonical
459+backwardCompose =
460+ Var.inCore ["Basics"] "<<"
461+
462+
463+type Assoc =
464+ Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
465+
466+
467+left :: Assoc
468+left =
469+ flip (,)
470+
471+
472+right :: Assoc
473+right =
474+ (,)
475+
476+
477+collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
478+collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
479+ case expr of
480+ Expr.Binop op leftExpr rightExpr | op == desiredOp ->
481+ let
482+ (func, arg) = assoc leftExpr rightExpr
483+ in
484+ A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
485+
486+ _ ->
487+ makeRoot annExpr