git.iko.soy/elk
on 684504e help sign in
elk/src/Optimize/DecisionTree.hs
467 lines 11 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.DecisionTree where
3{- To learn more about how this works, definitely read through:
5 "When Do Match-Compilation Heuristics Matter?"
7by Kevin Scott and Norman Ramsey. The rough idea is that we start with a simple
8list of patterns and expressions, and then turn that into a "decision tree"
9that requires as few tests as possible to make it to a leaf. Read the paper, it
10explains this extraordinarily well! We are currently using the same heuristics
11as SML/NJ to get nice trees.
12-}
14import Control.Arrow (second)
15import qualified Data.List as List
16import qualified Data.Map as Map
17import qualified Data.Set as Set
18import qualified Data.Maybe as Maybe
20import qualified AST.Helpers as Help
21import qualified AST.Literal as L
22import qualified AST.Pattern as P
23import qualified AST.Variable as Var
24import qualified Reporting.Annotation as A
27type CPattern = P.Canonical
30-- COMPILE CASES
32{-| Users of this module will mainly interact with this function. It takes
33some normal branches and gives out a decision tree that has "labels" at all
34the leafs and a dictionary that maps these "labels" to the code that should
35run.
37If 2 or more leaves point to the same label, we need to do some tricks in JS to
38make that work nicely. When is JS getting goto?! ;) That is outside the scope
39of this module though.
40-}
41compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree
42compile variantDict rawBranches =
43 let
44 format (pattern, index) =
45 Branch index [(Empty, pattern)]
46 in
47 toDecisionTree variantDict (map format rawBranches)
50{-| When a certain union type is defined, you specify a certain number of tags.
51This helps us do a few optimizations:
53 * If there is only one possible tag, we can always skip checking what it is.
54 Tuples are a common example of this.
55 * If we use all possible tags, we can skip doing the last test. If we have
56 checked N-1 of N tags, there is no need to test the Nth, we know its the
57 one we want
59So this dictionary maps tags to the number of variants that exist, so it'll
60contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
61we can use for these optimizations.
62-}
63type VariantDict =
64 Map.Map Var.Home (Map.Map String Int)
67-- DECISION TREES
69data DecisionTree
70 = Match Int
71 | Decision
72 { _test :: Path
73 , _edges :: [(Test, DecisionTree)]
74 , _default :: Maybe DecisionTree
75 }
76 deriving (Eq)
79data Test
80 = Constructor Var.Canonical
81 | Literal L.Literal
82 deriving (Eq, Ord)
85data Path
86 = Position Int Path
87 | Field String Path
88 | Empty
89 | Alias
90 deriving (Eq)
93-- PATH HELPERS
95add :: Path -> Path -> Path
96add path finalLink =
97 case path of
98 Empty ->
99 finalLink
101 Alias ->
102 error "nothing should be added to an alias path"
104 Position index subpath ->
105 Position index (add subpath finalLink)
107 Field name subpath ->
108 Field name (add subpath finalLink)
111subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
112subPositions path patterns =
113 zipWith
114 (\index pattern -> (add path (Position index Empty), pattern))
115 [0..]
116 patterns
119-- ACTUALLY BUILD DECISION TREES
121data Branch =
122 Branch
123 { _goal :: Int
124 , _patterns :: [(Path, CPattern)]
125 }
128toDecisionTree :: VariantDict -> [Branch] -> DecisionTree
129toDecisionTree variantDict rawBranches =
130 let
131 branches =
132 map (flattenPatterns variantDict) rawBranches
133 in
134 case checkForMatch branches of
135 Just goal ->
136 Match goal
138 Nothing ->
139 let
140 path =
141 pickPath variantDict branches
143 (edges, fallback) =
144 gatherEdges variantDict branches path
146 decisionEdges =
147 map (second (toDecisionTree variantDict)) edges
148 in
149 case (decisionEdges, fallback) of
150 ([(_tag, decisionTree)], []) ->
151 decisionTree
153 (_, []) ->
154 Decision path decisionEdges Nothing
156 ([], _ : _) ->
157 toDecisionTree variantDict fallback
159 (_, _) ->
160 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
163isComplete :: VariantDict -> [Test] -> Bool
164isComplete variantDict tests =
165 case head tests of
166 Constructor var ->
167 getArity variantDict var == length tests
169 Literal (L.Boolean _) ->
170 length tests == 2
172 _ ->
173 False
176getArity :: VariantDict -> Var.Canonical -> Int
177getArity variantDict (Var.Canonical home name) =
178 case Map.lookup name =<< Map.lookup home variantDict of
179 Just arity ->
180 arity
182 Nothing ->
183 if Help.isTuple name then
184 read (drop 6 name)
186 else
187 error
188 "Since the Optimize phase happens after canonicalization and type \
189 \inference, it is impossible that a pattern cannot be found."
192-- FLATTEN PATTERNS
194{-| Flatten type aliases and use the VariantDict to figure out when a tag is
195the only variant so we can skip doing any tests on it.
196-}
197flattenPatterns :: VariantDict -> Branch -> Branch
198flattenPatterns variantDict (Branch goal pathPatterns) =
199 Branch goal (concatMap (flatten variantDict) pathPatterns)
202flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
203flatten variantDict pathPattern@(path, A.A ann pattern) =
204 case pattern of
205 P.Var _ ->
206 [pathPattern]
208 P.Anything ->
209 [pathPattern]
211 P.Alias alias realPattern ->
212 (add path Alias, A.A ann (P.Var alias))
213 :
214 flatten variantDict (path, realPattern)
216 P.Record _ ->
217 [pathPattern]
219 P.Data tag patterns ->
220 if getArity variantDict tag == 1 then
221 concatMap (flatten variantDict) (subPositions path patterns)
223 else
224 [pathPattern]
226 P.Literal _ ->
227 [pathPattern]
230-- SUCCESSFULLY MATCH
232{-| If the first branch has no more "decision points" we can finally take that
233path. If that is the case we give the resulting label and a mapping from free
234variables to "how to get their value". So a pattern like (Just (x,_)) will give
235us something like ("x" => value.0.0)
236-}
237checkForMatch :: [Branch] -> Maybe Int
238checkForMatch branches =
239 case branches of
240 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
241 Just goal
243 _ ->
244 Nothing
247-- GATHER OUTGOING EDGES
249gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
250gatherEdges variantDict branches path =
251 let
252 relevantTests =
253 testsAtPath path branches
255 allEdges =
256 map (edgesFor path branches) relevantTests
258 fallbacks =
259 if isComplete variantDict relevantTests then
260 []
261 else
262 filter (isIrrelevantTo path) branches
263 in
264 ( allEdges, fallbacks )
267-- FIND RELEVANT TESTS
269testsAtPath :: Path -> [Branch] -> [Test]
270testsAtPath selectedPath branches =
271 let
272 allTests =
273 Maybe.mapMaybe (testAtPath selectedPath) branches
275 skipVisited test curr@(uniqueTests, visitedTests) =
276 if Set.member test visitedTests then
277 curr
278 else
279 ( test : uniqueTests
280 , Set.insert test visitedTests
281 )
282 in
283 fst (foldr skipVisited ([], Set.empty) allTests)
286testAtPath :: Path -> Branch -> Maybe Test
287testAtPath selectedPath (Branch _ pathPatterns) =
288 case List.lookup selectedPath pathPatterns of
289 Nothing ->
290 Nothing
292 Just (A.A _ pattern) ->
293 case pattern of
294 P.Data name _ ->
295 Just (Constructor name)
297 P.Literal lit ->
298 Just (Literal lit)
300 P.Var _ ->
301 Nothing
303 P.Alias _ _ ->
304 error "aliases should never reach 'testAtPath' function"
306 P.Anything ->
307 Nothing
309 P.Record _ ->
310 Nothing
313-- BUILD EDGES
315edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
316edgesFor path branches test =
317 ( test
318 , Maybe.mapMaybe (toRelevantBranch test path) branches
319 )
322toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
323toRelevantBranch test path branch@(Branch goal pathPatterns) =
324 case extract path pathPatterns of
325 Just (start, A.A _ pattern, end) ->
326 case pattern of
327 P.Data name patterns ->
328 if test == Constructor name then
329 Just (Branch goal (start ++ subPositions path patterns ++ end))
331 else
332 Nothing
334 P.Literal lit ->
335 if test == Literal lit then
336 Just (Branch goal (start ++ end))
338 else
339 Nothing
341 _ ->
342 Just branch
344 _ ->
345 Just branch
348extract
349 :: Path
350 -> [(Path, CPattern)]
351 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
352extract selectedPath pathPatterns =
353 case pathPatterns of
354 [] ->
355 Nothing
357 first@(path, pattern) : rest ->
358 if path == selectedPath then
359 Just ([], pattern, rest)
361 else
362 case extract selectedPath rest of
363 Nothing ->
364 Nothing
366 Just (start, foundPattern, end) ->
367 Just (first : start, foundPattern, end)
370-- FIND IRRELEVANT BRANCHES
372isIrrelevantTo :: Path -> Branch -> Bool
373isIrrelevantTo selectedPath (Branch _ pathPatterns) =
374 case List.lookup selectedPath pathPatterns of
375 Nothing ->
376 True
378 Just pattern ->
379 not (needsTests pattern)
382needsTests :: CPattern -> Bool
383needsTests (A.A _ pattern) =
384 case pattern of
385 P.Var _ ->
386 False
388 P.Anything ->
389 False
391 P.Alias _ _ ->
392 error "aliases should never reach 'isIrrelevantTo' function"
394 P.Record _ ->
395 False
397 P.Data _ _ ->
398 True
400 P.Literal _ ->
401 True
404-- PICK A PATH
406pickPath :: VariantDict -> [Branch] -> Path
407pickPath variantDict branches =
408 let
409 allPaths =
410 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
411 in
412 case bests (addWeights (smallDefaults branches) allPaths) of
413 [path] ->
414 path
416 tiedPaths ->
417 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
420isChoicePath :: (Path, CPattern) -> Maybe Path
421isChoicePath (path, pattern) =
422 if needsTests pattern then
423 Just path
424 else
425 Nothing
428addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
429addWeights toWeight paths =
430 map (\path -> (path, toWeight path)) paths
433bests :: [(Path, Int)] -> [Path]
434bests allPaths =
435 case allPaths of
436 [] ->
437 error "Cannot choose the best of zero paths. This should never happen."
439 (headPath, headWeight) : weightedPaths ->
440 let
441 gatherMinimum acc@(minWeight, paths) (path, weight) =
442 if weight == minWeight then
443 (minWeight, path : paths)
445 else if weight < minWeight then
446 (weight, [path])
448 else
449 acc
450 in
451 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
454-- PATH PICKING HEURISTICS
456smallDefaults :: [Branch] -> Path -> Int
457smallDefaults branches path =
458 length (filter (isIrrelevantTo path) branches)
461smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
462smallBranchingFactor variantDict branches path =
463 let
464 (edges, fallback) =
465 gatherEdges variantDict branches path
466 in
467 length edges + (if null fallback then 0 else 1)