git.iko.soy/elk
on f4ea9e3 help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
492 lines 12 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 Jump
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 a
70 = Match a
71 | Decision
72 { _test :: Path
73 , _edges :: [(Test, DecisionTree a)]
74 , _default :: Maybe (DecisionTree a)
75 }
76 deriving (Eq)
79data Jump = Jump
80 { _result :: Int
81 , _substitutions :: [(String, Path)]
82 }
85data Test
86 = Constructor Var.Canonical
87 | Literal L.Literal
88 deriving (Eq, Ord)
91data Path
92 = Position Int Path
93 | Field String Path
94 | Empty
95 | Alias
96 deriving (Eq)
99-- PATH HELPERS
101add :: Path -> Path -> Path
102add path finalLink =
103 case path of
104 Empty ->
105 finalLink
107 Alias ->
108 error "nothing should be added to an alias path"
110 Position index subpath ->
111 Position index (add subpath finalLink)
113 Field name subpath ->
114 Field name (add subpath finalLink)
117subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
118subPositions path patterns =
119 zipWith
120 (\index pattern -> (add path (Position index Empty), pattern))
121 [0..]
122 patterns
125-- ACTUALLY BUILD DECISION TREES
127data Branch =
128 Branch
129 { _goal :: Int
130 , _patterns :: [(Path, CPattern)]
131 }
134toDecisionTree :: VariantDict -> [Branch] -> DecisionTree Jump
135toDecisionTree variantDict rawBranches =
136 let
137 branches =
138 map (flattenPatterns variantDict) rawBranches
139 in
140 case checkForMatch branches of
141 Just (goal, substitutions) ->
142 Match (Jump goal substitutions)
144 Nothing ->
145 let
146 path =
147 pickPath variantDict branches
149 (edges, fallback) =
150 gatherEdges variantDict branches path
152 decisionEdges =
153 map (second (toDecisionTree variantDict)) edges
154 in
155 case (decisionEdges, fallback) of
156 ([(_tag, decisionTree)], []) ->
157 decisionTree
159 (_, []) ->
160 Decision path decisionEdges Nothing
162 ([], _ : _) ->
163 toDecisionTree variantDict fallback
165 (_, _) ->
166 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
169isComplete :: VariantDict -> [Test] -> Bool
170isComplete variantDict tests =
171 case head tests of
172 Constructor var ->
173 getArity variantDict var == length tests
175 Literal (L.Boolean _) ->
176 length tests == 2
178 _ ->
179 False
182getArity :: VariantDict -> Var.Canonical -> Int
183getArity variantDict (Var.Canonical home name) =
184 case Map.lookup name =<< Map.lookup home variantDict of
185 Just arity ->
186 arity
188 Nothing ->
189 if Help.isTuple name then
190 read (drop 6 name)
192 else
193 error
194 "Since the Optimize phase happens after canonicalization and type \
195 \inference, it is impossible that a pattern cannot be found."
198-- FLATTEN PATTERNS
200{-| Flatten type aliases and use the VariantDict to figure out when a tag is
201the only variant so we can skip doing any tests on it.
202-}
203flattenPatterns :: VariantDict -> Branch -> Branch
204flattenPatterns variantDict (Branch goal pathPatterns) =
205 Branch goal (concatMap (flatten variantDict) pathPatterns)
208flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
209flatten variantDict pathPattern@(path, A.A ann pattern) =
210 case pattern of
211 P.Var _ ->
212 [pathPattern]
214 P.Anything ->
215 [pathPattern]
217 P.Alias alias realPattern ->
218 [ (add path Alias, A.A ann (P.Var alias))
219 , (path, realPattern)
220 ]
222 P.Record _ ->
223 [pathPattern]
225 P.Data tag patterns ->
226 if getArity variantDict tag == 1 then
227 concatMap (flatten variantDict) (subPositions path patterns)
229 else
230 [pathPattern]
232 P.Literal _ ->
233 [pathPattern]
236-- SUCCESSFULLY MATCH
238{-| If the first branch has no more "decision points" we can finally take that
239path. If that is the case we give the resulting label and a mapping from free
240variables to "how to get their value". So a pattern like (Just (x,_)) will give
241us something like ("x" => value.0.0)
242-}
243checkForMatch :: [Branch] -> Maybe (Int, [(String, Path)])
244checkForMatch branches =
245 case branches of
246 Branch goal patterns : _ ->
247 (,) goal `fmap` concat `fmap` mapM getSubstitution patterns
249 _ ->
250 Nothing
253getSubstitution :: (Path, CPattern) -> Maybe [(String, Path)]
254getSubstitution (path, A.A _ pattern) =
255 case pattern of
256 P.Var x ->
257 Just [(x, path)]
259 P.Anything ->
260 Just []
262 P.Alias _ _ ->
263 error "aliases should never reach 'getSubstitution' function"
265 P.Record fields ->
266 Just (map (\name -> (name, add path (Field name Empty))) fields)
268 P.Data _ _ ->
269 Nothing
271 P.Literal _ ->
272 Nothing
275-- GATHER OUTGOING EDGES
277gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
278gatherEdges variantDict branches path =
279 let
280 relevantTests =
281 testsAtPath path branches
283 allRawEdges =
284 map (edgesFor path branches) relevantTests
285 in
286 if isComplete variantDict relevantTests then
287 ( init allRawEdges
288 , snd (last allRawEdges)
289 )
291 else
292 ( allRawEdges
293 , filter (isIrrelevantTo path) branches
294 )
297-- FIND RELEVANT TESTS
299testsAtPath :: Path -> [Branch] -> [Test]
300testsAtPath selectedPath branches =
301 let
302 allTests =
303 Maybe.mapMaybe (testAtPath selectedPath) branches
305 skipVisited test curr@(uniqueTests, visitedTests) =
306 if Set.member test visitedTests then
307 curr
308 else
309 ( test : uniqueTests
310 , Set.insert test visitedTests
311 )
312 in
313 fst (foldr skipVisited ([], Set.empty) allTests)
316testAtPath :: Path -> Branch -> Maybe Test
317testAtPath selectedPath (Branch _ pathPatterns) =
318 case List.lookup selectedPath pathPatterns of
319 Nothing ->
320 Nothing
322 Just (A.A _ pattern) ->
323 case pattern of
324 P.Data name _ ->
325 Just (Constructor name)
327 P.Literal lit ->
328 Just (Literal lit)
330 P.Var _ ->
331 Nothing
333 P.Alias _ _ ->
334 error "aliases should never reach 'testAtPath' function"
336 P.Anything ->
337 Nothing
339 P.Record _ ->
340 Nothing
343-- BUILD EDGES
345edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
346edgesFor path branches test =
347 ( test
348 , Maybe.mapMaybe (toRelevantBranch test path) branches
349 )
352toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
353toRelevantBranch test path branch@(Branch goal pathPatterns) =
354 case extract path pathPatterns of
355 Just (start, A.A _ pattern, end) ->
356 case pattern of
357 P.Data name patterns ->
358 if test == Constructor name then
359 Just (Branch goal (start ++ subPositions path patterns ++ end))
361 else
362 Nothing
364 P.Literal lit ->
365 if test == Literal lit then
366 Just (Branch goal (start ++ end))
368 else
369 Nothing
371 _ ->
372 Just branch
374 _ ->
375 Just branch
378extract
379 :: Path
380 -> [(Path, CPattern)]
381 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
382extract selectedPath pathPatterns =
383 case pathPatterns of
384 [] ->
385 Nothing
387 first@(path, pattern) : rest ->
388 if path == selectedPath then
389 Just ([], pattern, rest)
391 else
392 case extract selectedPath rest of
393 Nothing ->
394 Nothing
396 Just (start, foundPattern, end) ->
397 Just (first : start, foundPattern, end)
400-- FIND IRRELEVANT BRANCHES
402isIrrelevantTo :: Path -> Branch -> Bool
403isIrrelevantTo selectedPath (Branch _ pathPatterns) =
404 case List.lookup selectedPath pathPatterns of
405 Nothing ->
406 True
408 Just pattern ->
409 not (needsTests pattern)
412needsTests :: CPattern -> Bool
413needsTests (A.A _ pattern) =
414 case pattern of
415 P.Var _ ->
416 False
418 P.Anything ->
419 False
421 P.Alias _ _ ->
422 error "aliases should never reach 'isIrrelevantTo' function"
424 P.Record _ ->
425 False
427 P.Data _ _ ->
428 True
430 P.Literal _ ->
431 True
434-- PICK A PATH
436pickPath :: VariantDict -> [Branch] -> Path
437pickPath variantDict branches =
438 let
439 allPaths =
440 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
441 in
442 case bests (addWeights (smallDefaults branches) allPaths) of
443 [path] ->
444 path
446 tiedPaths ->
447 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
450isChoicePath :: (Path, CPattern) -> Maybe Path
451isChoicePath (path, pattern) =
452 if needsTests pattern then
453 Just path
454 else
455 Nothing
458addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
459addWeights toWeight paths =
460 map (\path -> (path, toWeight path)) paths
463bests :: [(Path, Int)] -> [Path]
464bests ((path,weight):weightedPaths) =
465 let
466 gatherMinimum acc@(minWeight, paths) (path, weight) =
467 if weight == minWeight then
468 (minWeight, path : paths)
470 else if weight < minWeight then
471 (weight, [path])
473 else
474 acc
475 in
476 snd (List.foldl' gatherMinimum (weight, [path]) weightedPaths)
479-- PATH PICKING HEURISTICS
481smallDefaults :: [Branch] -> Path -> Int
482smallDefaults branches path =
483 length (filter (isIrrelevantTo path) branches)
486smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
487smallBranchingFactor variantDict branches path =
488 let
489 (edges, fallback) =
490 gatherEdges variantDict branches path
491 in
492 length edges + (if null fallback then 0 else 1)