git.iko.soy/elk
on befb3df help sign in
elk/src/Optimize/DecisionTree.hs
540 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 Data.Binary
16import qualified Data.List as List
17import qualified Data.Map as Map
18import qualified Data.Set as Set
19import qualified Data.Maybe as Maybe
20import qualified Data.Text as Text
21import Data.Text (Text)
23import qualified AST.Helpers as Help
24import qualified AST.Literal as L
25import qualified AST.Pattern as P
26import qualified AST.Variable as Var
27import qualified Reporting.Annotation as A
30type CPattern = P.Canonical
34-- COMPILE CASES
37{-| Users of this module will mainly interact with this function. It takes
38some normal branches and gives out a decision tree that has "labels" at all
39the leafs and a dictionary that maps these "labels" to the code that should
40run.
42If 2 or more leaves point to the same label, we need to do some tricks in JS to
43make that work nicely. When is JS getting goto?! ;) That is outside the scope
44of this module though.
45-}
46compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree
47compile variantDict rawBranches =
48 let
49 format (pattern, index) =
50 Branch index [(Empty, pattern)]
51 in
52 toDecisionTree variantDict (map format rawBranches)
55{-| When a certain union type is defined, you specify a certain number of tags.
56This helps us do a few optimizations:
58 * If there is only one possible tag, we can always skip checking what it is.
59 Tuples are a common example of this.
60 * If we use all possible tags, we can skip doing the last test. If we have
61 checked N-1 of N tags, there is no need to test the Nth, we know its the
62 one we want
64So this dictionary maps tags to the number of variants that exist, so it'll
65contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
66we can use for these optimizations.
67-}
68type VariantDict =
69 Map.Map Var.Home (Map.Map Text Int)
73-- DECISION TREES
76data DecisionTree
77 = Match Int
78 | Decision
79 { _test :: Path
80 , _edges :: [(Test, DecisionTree)]
81 , _default :: Maybe DecisionTree
82 }
83 deriving (Eq)
86data Test
87 = Constructor Var.Canonical
88 | Literal L.Literal
89 deriving (Eq, Ord)
92data Path
93 = Position Int Path
94 | Field Text Path
95 | Empty
96 | Alias
97 deriving (Eq)
101-- PATH HELPERS
104add :: Path -> Path -> Path
105add path finalLink =
106 case path of
107 Empty ->
108 finalLink
110 Alias ->
111 error "nothing should be added to an alias path"
113 Position index subpath ->
114 Position index (add subpath finalLink)
116 Field name subpath ->
117 Field name (add subpath finalLink)
120subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
121subPositions path patterns =
122 zipWith
123 (\index pattern -> (add path (Position index Empty), pattern))
124 [0..]
125 patterns
129-- ACTUALLY BUILD DECISION TREES
132data Branch =
133 Branch
134 { _goal :: Int
135 , _patterns :: [(Path, CPattern)]
136 }
139toDecisionTree :: VariantDict -> [Branch] -> DecisionTree
140toDecisionTree variantDict rawBranches =
141 let
142 branches =
143 map (flattenPatterns variantDict) rawBranches
144 in
145 case checkForMatch branches of
146 Just goal ->
147 Match goal
149 Nothing ->
150 let
151 path =
152 pickPath variantDict branches
154 (edges, fallback) =
155 gatherEdges variantDict branches path
157 decisionEdges =
158 map (second (toDecisionTree variantDict)) edges
159 in
160 case (decisionEdges, fallback) of
161 ([(_tag, decisionTree)], []) ->
162 decisionTree
164 (_, []) ->
165 Decision path decisionEdges Nothing
167 ([], _ : _) ->
168 toDecisionTree variantDict fallback
170 (_, _) ->
171 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
174isComplete :: VariantDict -> [Test] -> Bool
175isComplete variantDict tests =
176 case head tests of
177 Constructor var ->
178 getArity variantDict var == length tests
180 Literal (L.Boolean _) ->
181 length tests == 2
183 _ ->
184 False
187getArity :: VariantDict -> Var.Canonical -> Int
188getArity variantDict (Var.Canonical home name) =
189 case Map.lookup name =<< Map.lookup home variantDict of
190 Just arity ->
191 arity
193 Nothing ->
194 if Help.isTuple name then
195 read (drop 6 (Text.unpack name))
197 else
198 error
199 "Since the Optimize phase happens after canonicalization and type \
200 \inference, it is impossible that a pattern cannot be found."
204-- FLATTEN PATTERNS
207{-| Flatten type aliases and use the VariantDict to figure out when a tag is
208the only variant so we can skip doing any tests on it.
209-}
210flattenPatterns :: VariantDict -> Branch -> Branch
211flattenPatterns variantDict (Branch goal pathPatterns) =
212 Branch goal (concatMap (flatten variantDict) pathPatterns)
215flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
216flatten variantDict pathPattern@(path, A.A ann pattern) =
217 case pattern of
218 P.Var _ ->
219 [pathPattern]
221 P.Anything ->
222 [pathPattern]
224 P.Alias alias realPattern ->
225 (add path Alias, A.A ann (P.Var alias))
226 :
227 flatten variantDict (path, realPattern)
229 P.Record _ ->
230 [pathPattern]
232 P.Ctor tag patterns ->
233 if getArity variantDict tag == 1 then
234 concatMap (flatten variantDict) (subPositions path patterns)
236 else
237 [pathPattern]
239 P.Literal _ ->
240 [pathPattern]
244-- SUCCESSFULLY MATCH
247{-| If the first branch has no more "decision points" we can finally take that
248path. If that is the case we give the resulting label and a mapping from free
249variables to "how to get their value". So a pattern like (Just (x,_)) will give
250us something like ("x" => value.0.0)
251-}
252checkForMatch :: [Branch] -> Maybe Int
253checkForMatch branches =
254 case branches of
255 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
256 Just goal
258 _ ->
259 Nothing
263-- GATHER OUTGOING EDGES
266gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
267gatherEdges variantDict branches path =
268 let
269 relevantTests =
270 testsAtPath path branches
272 allEdges =
273 map (edgesFor path branches) relevantTests
275 fallbacks =
276 if isComplete variantDict relevantTests then
277 []
278 else
279 filter (isIrrelevantTo path) branches
280 in
281 ( allEdges, fallbacks )
285-- FIND RELEVANT TESTS
288testsAtPath :: Path -> [Branch] -> [Test]
289testsAtPath selectedPath branches =
290 let
291 allTests =
292 Maybe.mapMaybe (testAtPath selectedPath) branches
294 skipVisited test curr@(uniqueTests, visitedTests) =
295 if Set.member test visitedTests then
296 curr
297 else
298 ( test : uniqueTests
299 , Set.insert test visitedTests
300 )
301 in
302 fst (foldr skipVisited ([], Set.empty) allTests)
305testAtPath :: Path -> Branch -> Maybe Test
306testAtPath selectedPath (Branch _ pathPatterns) =
307 case List.lookup selectedPath pathPatterns of
308 Nothing ->
309 Nothing
311 Just (A.A _ pattern) ->
312 case pattern of
313 P.Ctor name _ ->
314 Just (Constructor name)
316 P.Literal lit ->
317 Just (Literal lit)
319 P.Var _ ->
320 Nothing
322 P.Alias _ _ ->
323 error "aliases should never reach 'testAtPath' function"
325 P.Anything ->
326 Nothing
328 P.Record _ ->
329 Nothing
333-- BUILD EDGES
336edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
337edgesFor path branches test =
338 ( test
339 , Maybe.mapMaybe (toRelevantBranch test path) branches
340 )
343toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
344toRelevantBranch test path branch@(Branch goal pathPatterns) =
345 case extract path pathPatterns of
346 Just (start, A.A _ pattern, end) ->
347 case pattern of
348 P.Ctor name patterns ->
349 if test == Constructor name then
350 Just (Branch goal (start ++ subPositions path patterns ++ end))
352 else
353 Nothing
355 P.Literal lit ->
356 if test == Literal lit then
357 Just (Branch goal (start ++ end))
359 else
360 Nothing
362 _ ->
363 Just branch
365 _ ->
366 Just branch
369extract
370 :: Path
371 -> [(Path, CPattern)]
372 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
373extract selectedPath pathPatterns =
374 case pathPatterns of
375 [] ->
376 Nothing
378 first@(path, pattern) : rest ->
379 if path == selectedPath then
380 Just ([], pattern, rest)
382 else
383 case extract selectedPath rest of
384 Nothing ->
385 Nothing
387 Just (start, foundPattern, end) ->
388 Just (first : start, foundPattern, end)
392-- FIND IRRELEVANT BRANCHES
395isIrrelevantTo :: Path -> Branch -> Bool
396isIrrelevantTo selectedPath (Branch _ pathPatterns) =
397 case List.lookup selectedPath pathPatterns of
398 Nothing ->
399 True
401 Just pattern ->
402 not (needsTests pattern)
405needsTests :: CPattern -> Bool
406needsTests (A.A _ pattern) =
407 case pattern of
408 P.Var _ ->
409 False
411 P.Anything ->
412 False
414 P.Alias _ _ ->
415 error "aliases should never reach 'isIrrelevantTo' function"
417 P.Record _ ->
418 False
420 P.Ctor _ _ ->
421 True
423 P.Literal _ ->
424 True
428-- PICK A PATH
431pickPath :: VariantDict -> [Branch] -> Path
432pickPath variantDict branches =
433 let
434 allPaths =
435 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
436 in
437 case bests (addWeights (smallDefaults branches) allPaths) of
438 [path] ->
439 path
441 tiedPaths ->
442 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
445isChoicePath :: (Path, CPattern) -> Maybe Path
446isChoicePath (path, pattern) =
447 if needsTests pattern then
448 Just path
449 else
450 Nothing
453addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
454addWeights toWeight paths =
455 map (\path -> (path, toWeight path)) paths
458bests :: [(Path, Int)] -> [Path]
459bests allPaths =
460 case allPaths of
461 [] ->
462 error "Cannot choose the best of zero paths. This should never happen."
464 (headPath, headWeight) : 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 (headWeight, [headPath]) weightedPaths)
480-- PATH PICKING HEURISTICS
483smallDefaults :: [Branch] -> Path -> Int
484smallDefaults branches path =
485 length (filter (isIrrelevantTo path) branches)
488smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
489smallBranchingFactor variantDict branches path =
490 let
491 (edges, fallback) =
492 gatherEdges variantDict branches path
493 in
494 length edges + (if null fallback then 0 else 1)
498-- BINARY
501instance Binary Test where
502 get =
503 do word <- getWord8
504 case word of
505 0 -> Constructor <$> get
506 1 -> Literal <$> get
507 _ -> error "problem getting DecisionTree.Test binary"
509 put test =
510 case test of
511 Constructor var ->
512 putWord8 0 >> put var
514 Literal literal ->
515 putWord8 1 >> put literal
518instance Binary Path where
519 get =
520 do word <- getWord8
521 case word of
522 0 -> Position <$> get <*> get
523 1 -> Field <$> get <*> get
524 2 -> pure Empty
525 3 -> pure Alias
526 _ -> error "problem getting DecisionTree.Path binary"
528 put path =
529 case path of
530 Position index rest ->
531 putWord8 0 >> put index >> put rest
533 Field field rest ->
534 putWord8 1 >> put field >> put rest
536 Empty ->
537 putWord8 2
539 Alias ->
540 putWord8 3