git.iko.soy/elk
on fe7bbbd help sign in
elk/src/Optimize/DecisionTree.hs
493 lines 12 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
19import qualified Data.Text as Text
20import Data.Text (Text)
22import qualified AST.Helpers as Help
23import qualified AST.Literal as L
24import qualified AST.Pattern as P
25import qualified AST.Variable as Var
26import qualified Reporting.Annotation as A
29type CPattern = P.Canonical
33-- COMPILE CASES
36{-| Users of this module will mainly interact with this function. It takes
37some normal branches and gives out a decision tree that has "labels" at all
38the leafs and a dictionary that maps these "labels" to the code that should
39run.
41If 2 or more leaves point to the same label, we need to do some tricks in JS to
42make that work nicely. When is JS getting goto?! ;) That is outside the scope
43of this module though.
44-}
45compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree
46compile variantDict rawBranches =
47 let
48 format (pattern, index) =
49 Branch index [(Empty, pattern)]
50 in
51 toDecisionTree variantDict (map format rawBranches)
54{-| When a certain union type is defined, you specify a certain number of tags.
55This helps us do a few optimizations:
57 * If there is only one possible tag, we can always skip checking what it is.
58 Tuples are a common example of this.
59 * If we use all possible tags, we can skip doing the last test. If we have
60 checked N-1 of N tags, there is no need to test the Nth, we know its the
61 one we want
63So this dictionary maps tags to the number of variants that exist, so it'll
64contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
65we can use for these optimizations.
66-}
67type VariantDict =
68 Map.Map Var.Home (Map.Map Text Int)
72-- DECISION TREES
75data DecisionTree
76 = Match Int
77 | Decision
78 { _test :: Path
79 , _edges :: [(Test, DecisionTree)]
80 , _default :: Maybe DecisionTree
81 }
82 deriving (Eq)
85data Test
86 = Constructor Var.Canonical
87 | Literal L.Literal
88 deriving (Eq, Ord)
91data Path
92 = Position Int Path
93 | Field Text Path
94 | Empty
95 | Alias
96 deriving (Eq)
100-- PATH HELPERS
103add :: Path -> Path -> Path
104add path finalLink =
105 case path of
106 Empty ->
107 finalLink
109 Alias ->
110 error "nothing should be added to an alias path"
112 Position index subpath ->
113 Position index (add subpath finalLink)
115 Field name subpath ->
116 Field name (add subpath finalLink)
119subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
120subPositions path patterns =
121 zipWith
122 (\index pattern -> (add path (Position index Empty), pattern))
123 [0..]
124 patterns
128-- ACTUALLY BUILD DECISION TREES
131data Branch =
132 Branch
133 { _goal :: Int
134 , _patterns :: [(Path, CPattern)]
135 }
138toDecisionTree :: VariantDict -> [Branch] -> DecisionTree
139toDecisionTree variantDict rawBranches =
140 let
141 branches =
142 map (flattenPatterns variantDict) rawBranches
143 in
144 case checkForMatch branches of
145 Just goal ->
146 Match goal
148 Nothing ->
149 let
150 path =
151 pickPath variantDict branches
153 (edges, fallback) =
154 gatherEdges variantDict branches path
156 decisionEdges =
157 map (second (toDecisionTree variantDict)) edges
158 in
159 case (decisionEdges, fallback) of
160 ([(_tag, decisionTree)], []) ->
161 decisionTree
163 (_, []) ->
164 Decision path decisionEdges Nothing
166 ([], _ : _) ->
167 toDecisionTree variantDict fallback
169 (_, _) ->
170 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
173isComplete :: VariantDict -> [Test] -> Bool
174isComplete variantDict tests =
175 case head tests of
176 Constructor var ->
177 getArity variantDict var == length tests
179 Literal (L.Boolean _) ->
180 length tests == 2
182 _ ->
183 False
186getArity :: VariantDict -> Var.Canonical -> Int
187getArity variantDict (Var.Canonical home name) =
188 case Map.lookup name =<< Map.lookup home variantDict of
189 Just arity ->
190 arity
192 Nothing ->
193 if Help.isTuple name then
194 read (drop 6 (Text.unpack name))
196 else
197 error
198 "Since the Optimize phase happens after canonicalization and type \
199 \inference, it is impossible that a pattern cannot be found."
203-- FLATTEN PATTERNS
206{-| Flatten type aliases and use the VariantDict to figure out when a tag is
207the only variant so we can skip doing any tests on it.
208-}
209flattenPatterns :: VariantDict -> Branch -> Branch
210flattenPatterns variantDict (Branch goal pathPatterns) =
211 Branch goal (concatMap (flatten variantDict) pathPatterns)
214flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
215flatten variantDict pathPattern@(path, A.A ann pattern) =
216 case pattern of
217 P.Var _ ->
218 [pathPattern]
220 P.Anything ->
221 [pathPattern]
223 P.Alias alias realPattern ->
224 (add path Alias, A.A ann (P.Var alias))
225 :
226 flatten variantDict (path, realPattern)
228 P.Record _ ->
229 [pathPattern]
231 P.Ctor tag patterns ->
232 if getArity variantDict tag == 1 then
233 concatMap (flatten variantDict) (subPositions path patterns)
235 else
236 [pathPattern]
238 P.Literal _ ->
239 [pathPattern]
243-- SUCCESSFULLY MATCH
246{-| If the first branch has no more "decision points" we can finally take that
247path. If that is the case we give the resulting label and a mapping from free
248variables to "how to get their value". So a pattern like (Just (x,_)) will give
249us something like ("x" => value.0.0)
250-}
251checkForMatch :: [Branch] -> Maybe Int
252checkForMatch branches =
253 case branches of
254 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
255 Just goal
257 _ ->
258 Nothing
262-- GATHER OUTGOING EDGES
265gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
266gatherEdges variantDict branches path =
267 let
268 relevantTests =
269 testsAtPath path branches
271 allEdges =
272 map (edgesFor path branches) relevantTests
274 fallbacks =
275 if isComplete variantDict relevantTests then
276 []
277 else
278 filter (isIrrelevantTo path) branches
279 in
280 ( allEdges, fallbacks )
284-- FIND RELEVANT TESTS
287testsAtPath :: Path -> [Branch] -> [Test]
288testsAtPath selectedPath branches =
289 let
290 allTests =
291 Maybe.mapMaybe (testAtPath selectedPath) branches
293 skipVisited test curr@(uniqueTests, visitedTests) =
294 if Set.member test visitedTests then
295 curr
296 else
297 ( test : uniqueTests
298 , Set.insert test visitedTests
299 )
300 in
301 fst (foldr skipVisited ([], Set.empty) allTests)
304testAtPath :: Path -> Branch -> Maybe Test
305testAtPath selectedPath (Branch _ pathPatterns) =
306 case List.lookup selectedPath pathPatterns of
307 Nothing ->
308 Nothing
310 Just (A.A _ pattern) ->
311 case pattern of
312 P.Ctor name _ ->
313 Just (Constructor name)
315 P.Literal lit ->
316 Just (Literal lit)
318 P.Var _ ->
319 Nothing
321 P.Alias _ _ ->
322 error "aliases should never reach 'testAtPath' function"
324 P.Anything ->
325 Nothing
327 P.Record _ ->
328 Nothing
332-- BUILD EDGES
335edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
336edgesFor path branches test =
337 ( test
338 , Maybe.mapMaybe (toRelevantBranch test path) branches
339 )
342toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
343toRelevantBranch test path branch@(Branch goal pathPatterns) =
344 case extract path pathPatterns of
345 Just (start, A.A _ pattern, end) ->
346 case pattern of
347 P.Ctor name patterns ->
348 if test == Constructor name then
349 Just (Branch goal (start ++ subPositions path patterns ++ end))
351 else
352 Nothing
354 P.Literal lit ->
355 if test == Literal lit then
356 Just (Branch goal (start ++ end))
358 else
359 Nothing
361 _ ->
362 Just branch
364 _ ->
365 Just branch
368extract
369 :: Path
370 -> [(Path, CPattern)]
371 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
372extract selectedPath pathPatterns =
373 case pathPatterns of
374 [] ->
375 Nothing
377 first@(path, pattern) : rest ->
378 if path == selectedPath then
379 Just ([], pattern, rest)
381 else
382 case extract selectedPath rest of
383 Nothing ->
384 Nothing
386 Just (start, foundPattern, end) ->
387 Just (first : start, foundPattern, end)
391-- FIND IRRELEVANT BRANCHES
394isIrrelevantTo :: Path -> Branch -> Bool
395isIrrelevantTo selectedPath (Branch _ pathPatterns) =
396 case List.lookup selectedPath pathPatterns of
397 Nothing ->
398 True
400 Just pattern ->
401 not (needsTests pattern)
404needsTests :: CPattern -> Bool
405needsTests (A.A _ pattern) =
406 case pattern of
407 P.Var _ ->
408 False
410 P.Anything ->
411 False
413 P.Alias _ _ ->
414 error "aliases should never reach 'isIrrelevantTo' function"
416 P.Record _ ->
417 False
419 P.Ctor _ _ ->
420 True
422 P.Literal _ ->
423 True
427-- PICK A PATH
430pickPath :: VariantDict -> [Branch] -> Path
431pickPath variantDict branches =
432 let
433 allPaths =
434 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
435 in
436 case bests (addWeights (smallDefaults branches) allPaths) of
437 [path] ->
438 path
440 tiedPaths ->
441 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
444isChoicePath :: (Path, CPattern) -> Maybe Path
445isChoicePath (path, pattern) =
446 if needsTests pattern then
447 Just path
448 else
449 Nothing
452addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
453addWeights toWeight paths =
454 map (\path -> (path, toWeight path)) paths
457bests :: [(Path, Int)] -> [Path]
458bests allPaths =
459 case allPaths of
460 [] ->
461 error "Cannot choose the best of zero paths. This should never happen."
463 (headPath, headWeight) : weightedPaths ->
464 let
465 gatherMinimum acc@(minWeight, paths) (path, weight) =
466 if weight == minWeight then
467 (minWeight, path : paths)
469 else if weight < minWeight then
470 (weight, [path])
472 else
473 acc
474 in
475 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
479-- PATH PICKING HEURISTICS
482smallDefaults :: [Branch] -> Path -> Int
483smallDefaults branches path =
484 length (filter (isIrrelevantTo path) branches)
487smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
488smallBranchingFactor variantDict branches path =
489 let
490 (edges, fallback) =
491 gatherEdges variantDict branches path
492 in
493 length edges + (if null fallback then 0 else 1)