git.iko.soy/elk
on 8ac06ce help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
489 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 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
166isComplete :: VariantDict -> [Test] -> Bool
167isComplete variantDict tests =
168 case head tests of
169 Constructor var ->
170 getArity variantDict var == length tests
172 Literal (L.Boolean _) ->
173 length tests == 2
175 _ ->
176 False
179getArity :: VariantDict -> Var.Canonical -> Int
180getArity variantDict (Var.Canonical home name) =
181 case Map.lookup name =<< Map.lookup home variantDict of
182 Just arity ->
183 arity
185 Nothing ->
186 if Help.isTuple name then
187 read (drop 6 name)
189 else
190 error
191 "Since the Optimize phase happens after canonicalization and type \
192 \inference, it is impossible that a pattern cannot be found."
195-- FLATTEN PATTERNS
197{-| Flatten type aliases and use the VariantDict to figure out when a tag is
198the only variant so we can skip doing any tests on it.
199-}
200flattenPatterns :: VariantDict -> Branch -> Branch
201flattenPatterns variantDict (Branch goal pathPatterns) =
202 Branch goal (concatMap (flatten variantDict) pathPatterns)
205flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
206flatten variantDict pathPattern@(path, A.A ann pattern) =
207 case pattern of
208 P.Var _ ->
209 [pathPattern]
211 P.Anything ->
212 [pathPattern]
214 P.Alias alias realPattern ->
215 [ (add path Alias, A.A ann (P.Var alias))
216 , (path, realPattern)
217 ]
219 P.Record _ ->
220 [pathPattern]
222 P.Data tag patterns ->
223 if getArity variantDict tag == 1 then
224 concatMap (flatten variantDict) (subPositions path patterns)
226 else
227 [pathPattern]
229 P.Literal _ ->
230 [pathPattern]
233-- SUCCESSFULLY MATCH
235{-| If the first branch has no more "decision points" we can finally take that
236path. If that is the case we give the resulting label and a mapping from free
237variables to "how to get their value". So a pattern like (Just (x,_)) will give
238us something like ("x" => value.0.0)
239-}
240checkForMatch :: [Branch] -> Maybe (Int, [(String, Path)])
241checkForMatch branches =
242 case branches of
243 Branch goal patterns : _ ->
244 (,) goal `fmap` concat `fmap` mapM getSubstitution patterns
246 _ ->
247 Nothing
250getSubstitution :: (Path, CPattern) -> Maybe [(String, Path)]
251getSubstitution (path, A.A _ pattern) =
252 case pattern of
253 P.Var x ->
254 Just [(x, path)]
256 P.Anything ->
257 Just []
259 P.Alias _ _ ->
260 error "aliases should never reach 'getSubstitution' function"
262 P.Record fields ->
263 Just (map (\name -> (name, add path (Field name Empty))) fields)
265 P.Data _ _ ->
266 Nothing
268 P.Literal _ ->
269 Nothing
272-- GATHER OUTGOING EDGES
274gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
275gatherEdges variantDict branches path =
276 let
277 relevantTests =
278 testsAtPath path branches
280 allRawEdges =
281 map (edgesFor path branches) relevantTests
282 in
283 if isComplete variantDict relevantTests then
284 ( init allRawEdges
285 , snd (last allRawEdges)
286 )
288 else
289 ( allRawEdges
290 , filter (isIrrelevantTo path) branches
291 )
294-- FIND RELEVANT TESTS
296testsAtPath :: Path -> [Branch] -> [Test]
297testsAtPath selectedPath branches =
298 let
299 allTests =
300 Maybe.mapMaybe (testAtPath selectedPath) branches
302 skipVisited test curr@(uniqueTests, visitedTests) =
303 if Set.member test visitedTests then
304 curr
305 else
306 ( test : uniqueTests
307 , Set.insert test visitedTests
308 )
309 in
310 fst (foldr skipVisited ([], Set.empty) allTests)
313testAtPath :: Path -> Branch -> Maybe Test
314testAtPath selectedPath (Branch _ pathPatterns) =
315 case List.lookup selectedPath pathPatterns of
316 Nothing ->
317 Nothing
319 Just (A.A _ pattern) ->
320 case pattern of
321 P.Data name _ ->
322 Just (Constructor name)
324 P.Literal lit ->
325 Just (Literal lit)
327 P.Var _ ->
328 Nothing
330 P.Alias _ _ ->
331 error "aliases should never reach 'testAtPath' function"
333 P.Anything ->
334 Nothing
336 P.Record _ ->
337 Nothing
340-- BUILD EDGES
342edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
343edgesFor path branches test =
344 ( test
345 , Maybe.mapMaybe (toRelevantBranch test path) branches
346 )
349toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
350toRelevantBranch test path branch@(Branch goal pathPatterns) =
351 case extract path pathPatterns of
352 Just (start, A.A _ pattern, end) ->
353 case pattern of
354 P.Data name patterns ->
355 if test == Constructor name then
356 Just (Branch goal (start ++ subPositions path patterns ++ end))
358 else
359 Nothing
361 P.Literal lit ->
362 if test == Literal lit then
363 Just (Branch goal (start ++ end))
365 else
366 Nothing
368 _ ->
369 Just branch
371 _ ->
372 Just branch
375extract
376 :: Path
377 -> [(Path, CPattern)]
378 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
379extract selectedPath pathPatterns =
380 case pathPatterns of
381 [] ->
382 Nothing
384 first@(path, pattern) : rest ->
385 if path == selectedPath then
386 Just ([], pattern, rest)
388 else
389 case extract selectedPath rest of
390 Nothing ->
391 Nothing
393 Just (start, foundPattern, end) ->
394 Just (first : start, foundPattern, end)
397-- FIND IRRELEVANT BRANCHES
399isIrrelevantTo :: Path -> Branch -> Bool
400isIrrelevantTo selectedPath (Branch _ pathPatterns) =
401 case List.lookup selectedPath pathPatterns of
402 Nothing ->
403 True
405 Just pattern ->
406 not (needsTests pattern)
409needsTests :: CPattern -> Bool
410needsTests (A.A _ pattern) =
411 case pattern of
412 P.Var _ ->
413 False
415 P.Anything ->
416 False
418 P.Alias _ _ ->
419 error "aliases should never reach 'isIrrelevantTo' function"
421 P.Record _ ->
422 False
424 P.Data _ _ ->
425 True
427 P.Literal _ ->
428 True
431-- PICK A PATH
433pickPath :: VariantDict -> [Branch] -> Path
434pickPath variantDict branches =
435 let
436 allPaths =
437 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
438 in
439 case bests (addWeights (smallDefaults branches) allPaths) of
440 [path] ->
441 path
443 tiedPaths ->
444 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
447isChoicePath :: (Path, CPattern) -> Maybe Path
448isChoicePath (path, pattern) =
449 if needsTests pattern then
450 Just path
451 else
452 Nothing
455addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
456addWeights toWeight paths =
457 map (\path -> (path, toWeight path)) paths
460bests :: [(Path, Int)] -> [Path]
461bests ((path,weight):weightedPaths) =
462 let
463 gatherMinimum acc@(minWeight, paths) (path, weight) =
464 if weight == minWeight then
465 (minWeight, path : paths)
467 else if weight < minWeight then
468 (weight, [path])
470 else
471 acc
472 in
473 snd (List.foldl' gatherMinimum (weight, [path]) weightedPaths)
476-- PATH PICKING HEURISTICS
478smallDefaults :: [Branch] -> Path -> Int
479smallDefaults branches path =
480 length (filter (isIrrelevantTo path) branches)
483smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
484smallBranchingFactor variantDict branches path =
485 let
486 (edges, fallback) =
487 gatherEdges variantDict branches path
488 in
489 length edges + (if null fallback then 0 else 1)