git.iko.soy/elk
on 176fac8 help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
464 lines 11 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Patterns.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.CanonicalPattern
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 , (path, realPattern)
214 ]
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 allRawEdges =
256 map (edgesFor path branches) relevantTests
257 in
258 if isComplete variantDict relevantTests then
259 ( init allRawEdges
260 , snd (last allRawEdges)
261 )
263 else
264 ( allRawEdges
265 , filter (isIrrelevantTo path) branches
266 )
269-- FIND RELEVANT TESTS
271testsAtPath :: Path -> [Branch] -> [Test]
272testsAtPath selectedPath branches =
273 let
274 allTests =
275 Maybe.mapMaybe (testAtPath selectedPath) branches
277 skipVisited test curr@(uniqueTests, visitedTests) =
278 if Set.member test visitedTests then
279 curr
280 else
281 ( test : uniqueTests
282 , Set.insert test visitedTests
283 )
284 in
285 fst (foldr skipVisited ([], Set.empty) allTests)
288testAtPath :: Path -> Branch -> Maybe Test
289testAtPath selectedPath (Branch _ pathPatterns) =
290 case List.lookup selectedPath pathPatterns of
291 Nothing ->
292 Nothing
294 Just (A.A _ pattern) ->
295 case pattern of
296 P.Data name _ ->
297 Just (Constructor name)
299 P.Literal lit ->
300 Just (Literal lit)
302 P.Var _ ->
303 Nothing
305 P.Alias _ _ ->
306 error "aliases should never reach 'testAtPath' function"
308 P.Anything ->
309 Nothing
311 P.Record _ ->
312 Nothing
315-- BUILD EDGES
317edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
318edgesFor path branches test =
319 ( test
320 , Maybe.mapMaybe (toRelevantBranch test path) branches
321 )
324toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
325toRelevantBranch test path branch@(Branch goal pathPatterns) =
326 case extract path pathPatterns of
327 Just (start, A.A _ pattern, end) ->
328 case pattern of
329 P.Data name patterns ->
330 if test == Constructor name then
331 Just (Branch goal (start ++ subPositions path patterns ++ end))
333 else
334 Nothing
336 P.Literal lit ->
337 if test == Literal lit then
338 Just (Branch goal (start ++ end))
340 else
341 Nothing
343 _ ->
344 Just branch
346 _ ->
347 Just branch
350extract
351 :: Path
352 -> [(Path, CPattern)]
353 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
354extract selectedPath pathPatterns =
355 case pathPatterns of
356 [] ->
357 Nothing
359 first@(path, pattern) : rest ->
360 if path == selectedPath then
361 Just ([], pattern, rest)
363 else
364 case extract selectedPath rest of
365 Nothing ->
366 Nothing
368 Just (start, foundPattern, end) ->
369 Just (first : start, foundPattern, end)
372-- FIND IRRELEVANT BRANCHES
374isIrrelevantTo :: Path -> Branch -> Bool
375isIrrelevantTo selectedPath (Branch _ pathPatterns) =
376 case List.lookup selectedPath pathPatterns of
377 Nothing ->
378 True
380 Just pattern ->
381 not (needsTests pattern)
384needsTests :: CPattern -> Bool
385needsTests (A.A _ pattern) =
386 case pattern of
387 P.Var _ ->
388 False
390 P.Anything ->
391 False
393 P.Alias _ _ ->
394 error "aliases should never reach 'isIrrelevantTo' function"
396 P.Record _ ->
397 False
399 P.Data _ _ ->
400 True
402 P.Literal _ ->
403 True
406-- PICK A PATH
408pickPath :: VariantDict -> [Branch] -> Path
409pickPath variantDict branches =
410 let
411 allPaths =
412 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
413 in
414 case bests (addWeights (smallDefaults branches) allPaths) of
415 [path] ->
416 path
418 tiedPaths ->
419 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
422isChoicePath :: (Path, CPattern) -> Maybe Path
423isChoicePath (path, pattern) =
424 if needsTests pattern then
425 Just path
426 else
427 Nothing
430addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
431addWeights toWeight paths =
432 map (\path -> (path, toWeight path)) paths
435bests :: [(Path, Int)] -> [Path]
436bests ((path,weight):weightedPaths) =
437 let
438 gatherMinimum acc@(minWeight, paths) (path, weight) =
439 if weight == minWeight then
440 (minWeight, path : paths)
442 else if weight < minWeight then
443 (weight, [path])
445 else
446 acc
447 in
448 snd (List.foldl' gatherMinimum (weight, [path]) weightedPaths)
451-- PATH PICKING HEURISTICS
453smallDefaults :: [Branch] -> Path -> Int
454smallDefaults branches path =
455 length (filter (isIrrelevantTo path) branches)
458smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
459smallBranchingFactor variantDict branches path =
460 let
461 (edges, fallback) =
462 gatherEdges variantDict branches path
463 in
464 length edges + (if null fallback then 0 else 1)