git.iko.soy/elk
on 294ee07 help sign in
elk/src/Optimize/DecisionTree.hs
469 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
32-- COMPILE CASES
34{-| Users of this module will mainly interact with this function. It takes
35some normal branches and gives out a decision tree that has "labels" at all
36the leafs and a dictionary that maps these "labels" to the code that should
37run.
39If 2 or more leaves point to the same label, we need to do some tricks in JS to
40make that work nicely. When is JS getting goto?! ;) That is outside the scope
41of this module though.
42-}
43compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree
44compile variantDict rawBranches =
45 let
46 format (pattern, index) =
47 Branch index [(Empty, pattern)]
48 in
49 toDecisionTree variantDict (map format rawBranches)
52{-| When a certain union type is defined, you specify a certain number of tags.
53This helps us do a few optimizations:
55 * If there is only one possible tag, we can always skip checking what it is.
56 Tuples are a common example of this.
57 * If we use all possible tags, we can skip doing the last test. If we have
58 checked N-1 of N tags, there is no need to test the Nth, we know its the
59 one we want
61So this dictionary maps tags to the number of variants that exist, so it'll
62contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
63we can use for these optimizations.
64-}
65type VariantDict =
66 Map.Map Var.Home (Map.Map Text Int)
69-- DECISION TREES
71data DecisionTree
72 = Match Int
73 | Decision
74 { _test :: Path
75 , _edges :: [(Test, DecisionTree)]
76 , _default :: Maybe DecisionTree
77 }
78 deriving (Eq)
81data Test
82 = Constructor Var.Canonical
83 | Literal L.Literal
84 deriving (Eq, Ord)
87data Path
88 = Position Int Path
89 | Field Text Path
90 | Empty
91 | Alias
92 deriving (Eq)
95-- PATH HELPERS
97add :: Path -> Path -> Path
98add path finalLink =
99 case path of
100 Empty ->
101 finalLink
103 Alias ->
104 error "nothing should be added to an alias path"
106 Position index subpath ->
107 Position index (add subpath finalLink)
109 Field name subpath ->
110 Field name (add subpath finalLink)
113subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
114subPositions path patterns =
115 zipWith
116 (\index pattern -> (add path (Position index Empty), pattern))
117 [0..]
118 patterns
121-- ACTUALLY BUILD DECISION TREES
123data Branch =
124 Branch
125 { _goal :: Int
126 , _patterns :: [(Path, CPattern)]
127 }
130toDecisionTree :: VariantDict -> [Branch] -> DecisionTree
131toDecisionTree variantDict rawBranches =
132 let
133 branches =
134 map (flattenPatterns variantDict) rawBranches
135 in
136 case checkForMatch branches of
137 Just goal ->
138 Match goal
140 Nothing ->
141 let
142 path =
143 pickPath variantDict branches
145 (edges, fallback) =
146 gatherEdges variantDict branches path
148 decisionEdges =
149 map (second (toDecisionTree variantDict)) edges
150 in
151 case (decisionEdges, fallback) of
152 ([(_tag, decisionTree)], []) ->
153 decisionTree
155 (_, []) ->
156 Decision path decisionEdges Nothing
158 ([], _ : _) ->
159 toDecisionTree variantDict fallback
161 (_, _) ->
162 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
165isComplete :: VariantDict -> [Test] -> Bool
166isComplete variantDict tests =
167 case head tests of
168 Constructor var ->
169 getArity variantDict var == length tests
171 Literal (L.Boolean _) ->
172 length tests == 2
174 _ ->
175 False
178getArity :: VariantDict -> Var.Canonical -> Int
179getArity variantDict (Var.Canonical home name) =
180 case Map.lookup name =<< Map.lookup home variantDict of
181 Just arity ->
182 arity
184 Nothing ->
185 if Help.isTuple name then
186 read (drop 6 (Text.unpack name))
188 else
189 error
190 "Since the Optimize phase happens after canonicalization and type \
191 \inference, it is impossible that a pattern cannot be found."
194-- FLATTEN PATTERNS
196{-| Flatten type aliases and use the VariantDict to figure out when a tag is
197the only variant so we can skip doing any tests on it.
198-}
199flattenPatterns :: VariantDict -> Branch -> Branch
200flattenPatterns variantDict (Branch goal pathPatterns) =
201 Branch goal (concatMap (flatten variantDict) pathPatterns)
204flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
205flatten variantDict pathPattern@(path, A.A ann pattern) =
206 case pattern of
207 P.Var _ ->
208 [pathPattern]
210 P.Anything ->
211 [pathPattern]
213 P.Alias alias realPattern ->
214 (add path Alias, A.A ann (P.Var alias))
215 :
216 flatten variantDict (path, realPattern)
218 P.Record _ ->
219 [pathPattern]
221 P.Ctor tag patterns ->
222 if getArity variantDict tag == 1 then
223 concatMap (flatten variantDict) (subPositions path patterns)
225 else
226 [pathPattern]
228 P.Literal _ ->
229 [pathPattern]
232-- SUCCESSFULLY MATCH
234{-| If the first branch has no more "decision points" we can finally take that
235path. If that is the case we give the resulting label and a mapping from free
236variables to "how to get their value". So a pattern like (Just (x,_)) will give
237us something like ("x" => value.0.0)
238-}
239checkForMatch :: [Branch] -> Maybe Int
240checkForMatch branches =
241 case branches of
242 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
243 Just goal
245 _ ->
246 Nothing
249-- GATHER OUTGOING EDGES
251gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
252gatherEdges variantDict branches path =
253 let
254 relevantTests =
255 testsAtPath path branches
257 allEdges =
258 map (edgesFor path branches) relevantTests
260 fallbacks =
261 if isComplete variantDict relevantTests then
262 []
263 else
264 filter (isIrrelevantTo path) branches
265 in
266 ( allEdges, fallbacks )
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.Ctor 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.Ctor 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.Ctor _ _ ->
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 allPaths =
437 case allPaths of
438 [] ->
439 error "Cannot choose the best of zero paths. This should never happen."
441 (headPath, headWeight) : weightedPaths ->
442 let
443 gatherMinimum acc@(minWeight, paths) (path, weight) =
444 if weight == minWeight then
445 (minWeight, path : paths)
447 else if weight < minWeight then
448 (weight, [path])
450 else
451 acc
452 in
453 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
456-- PATH PICKING HEURISTICS
458smallDefaults :: [Branch] -> Path -> Int
459smallDefaults branches path =
460 length (filter (isIrrelevantTo path) branches)
463smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
464smallBranchingFactor variantDict branches path =
465 let
466 (edges, fallback) =
467 gatherEdges variantDict branches path
468 in
469 length edges + (if null fallback then 0 else 1)