git.iko.soy/elk
on 4654cef help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
409 lines 9.9 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 Data.Function (on)
16import qualified Data.List as List
17import qualified Data.Map as Map
18import Data.Map ((!))
19import qualified Data.Maybe as Maybe
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.Canonical 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 }
78data Jump = Jump
79 { _result :: Int
80 , _substitutions :: [(String, Path)]
81 }
84data Test
85 = Constructor Var.Canonical
86 | Literal L.Literal
87 deriving (Eq, Ord)
90data Path
91 = Position Int Path
92 | Field String Path
93 | Empty
94 | Alias
95 deriving (Eq)
98-- PATH HELPERS
100add :: Path -> Path -> Path
101add path finalLink =
102 case path of
103 Empty ->
104 finalLink
106 Alias ->
107 error "nothing should be added to an alias path"
109 Position index subpath ->
110 Position index (add subpath finalLink)
112 Field name subpath ->
113 Field name (add subpath finalLink)
116subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
117subPositions path patterns =
118 zipWith
119 (\index pattern -> (add path (Position index Empty), pattern))
120 [0..]
121 patterns
124-- ACTUALLY BUILD DECISION TREES
126data Branch =
127 Branch
128 { _goal :: Int
129 , _patterns :: [(Path, CPattern)]
130 }
133toDecisionTree :: VariantDict -> [Branch] -> DecisionTree Jump
134toDecisionTree variantDict rawBranches =
135 let
136 branches =
137 map (flattenPatterns variantDict) rawBranches
138 in
139 case checkForMatch branches of
140 Just (goal, substitutions) ->
141 Match (Jump goal substitutions)
143 Nothing ->
144 let
145 path =
146 smallDefaults branches
148 relevantTests =
149 testsAtPath path branches
151 allRawEdges =
152 map (edgesFor path branches) relevantTests
154 (rawEdges, unexplored) =
155 if isComplete variantDict relevantTests then
156 ( init allRawEdges
157 , snd (last allRawEdges)
158 )
160 else
161 ( allRawEdges
162 , filter (isIrrelevantTo path) branches
163 )
165 edges =
166 map (second (toDecisionTree variantDict)) rawEdges
167 in
168 case (edges, unexplored) of
169 ([(_tag, decisionTree)], []) ->
170 decisionTree
172 (_, []) ->
173 Decision path edges Nothing
175 (_, _) ->
176 Decision path edges (Just (toDecisionTree variantDict unexplored))
179isComplete :: VariantDict -> [Test] -> Bool
180isComplete variantDict tests =
181 case head tests of
182 Constructor name ->
183 (variantDict ! name) == length tests
185 Literal (L.Boolean _) ->
186 length tests == 2
188 _ ->
189 False
193-- FLATTEN PATTERNS
195{-| Flatten type aliases and use the VariantDict to figure out when a tag is
196the only variant so we can skip doing any tests on it.
197-}
198flattenPatterns :: VariantDict -> Branch -> Branch
199flattenPatterns variantDict (Branch goal pathPatterns) =
200 Branch goal (concatMap (flatten variantDict) pathPatterns)
203flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
204flatten variantDict pathPattern@(path, A.A ann pattern) =
205 case pattern of
206 P.Var _ ->
207 [pathPattern]
209 P.Anything ->
210 [pathPattern]
212 P.Alias alias realPattern ->
213 [ (add path Alias, A.A ann (P.Var alias))
214 , (path, realPattern)
215 ]
217 P.Record _ ->
218 [pathPattern]
220 P.Data tag patterns ->
221 if variantDict ! tag == 1 then
222 concatMap (flatten variantDict) (subPositions path patterns)
224 else
225 [pathPattern]
227 P.Literal _ ->
228 [pathPattern]
231-- SUCCESSFULLY MATCH
233{-| If the first branch has no more "decision points" we can finally take that
234path. If that is the case we give the resulting label and a mapping from free
235variables to "how to get their value". So a pattern like (Just (x,_)) will give
236us something like ("x" => value.0.0)
237-}
238checkForMatch :: [Branch] -> Maybe (Int, [(String, Path)])
239checkForMatch branches =
240 case branches of
241 Branch goal patterns : _ ->
242 (,) goal `fmap` concat `fmap` mapM getSubstitution patterns
244 _ ->
245 Nothing
248getSubstitution :: (Path, CPattern) -> Maybe [(String, Path)]
249getSubstitution (path, A.A _ pattern) =
250 case pattern of
251 P.Var x ->
252 Just [(x, path)]
254 P.Anything ->
255 Just []
257 P.Alias _ _ ->
258 error "aliases should never reach 'getSubstitution' function"
260 P.Record fields ->
261 Just (map (\name -> (name, add path (Field name Empty))) fields)
263 P.Data _ _ ->
264 Nothing
266 P.Literal _ ->
267 Nothing
270-- FIND RELEVANT TESTS
272testsAtPath :: Path -> [Branch] -> [Test]
273testsAtPath selectedPath branches =
274 List.nub $ Maybe.mapMaybe (testAtPath selectedPath) branches
277testAtPath :: Path -> Branch -> Maybe Test
278testAtPath selectedPath (Branch _ pathPatterns) =
279 case List.lookup selectedPath pathPatterns of
280 Nothing ->
281 Nothing
283 Just (A.A _ pattern) ->
284 case pattern of
285 P.Data name _ ->
286 Just (Constructor name)
288 P.Literal lit ->
289 Just (Literal lit)
291 P.Var _ ->
292 Nothing
294 P.Alias _ _ ->
295 error "aliases should never reach 'testAtPath' function"
297 P.Anything ->
298 Nothing
300 P.Record _ ->
301 Nothing
304-- BUILD EDGES
306edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
307edgesFor path branches test =
308 ( test
309 , Maybe.mapMaybe (toRelevantBranch test path) branches
310 )
313toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
314toRelevantBranch test path branch@(Branch goal pathPatterns) =
315 case extract path pathPatterns of
316 Just (start, A.A _ pattern, end) ->
317 case pattern of
318 P.Data name patterns | test == Constructor name ->
319 Just (Branch goal (start ++ subPositions path patterns ++ end))
321 P.Literal lit | test == Literal lit ->
322 Just (Branch goal (start ++ end))
324 _ ->
325 Nothing
327 _ ->
328 Just branch
331extract
332 :: Path
333 -> [(Path, CPattern)]
334 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
335extract selectedPath pathPatterns =
336 case pathPatterns of
337 [] ->
338 Nothing
340 first@(path, pattern) : rest ->
341 if path == selectedPath then
342 Just ([], pattern, rest)
344 else
345 case extract selectedPath rest of
346 Nothing ->
347 Nothing
349 Just (start, foundPattern, end) ->
350 Just (first : start, foundPattern, end)
353-- FIND IRRELEVANT BRANCHES
355isIrrelevantTo :: Path -> Branch -> Bool
356isIrrelevantTo selectedPath (Branch _ pathPatterns) =
357 case List.lookup selectedPath pathPatterns of
358 Nothing ->
359 True
361 Just pattern ->
362 not (needsTests pattern)
365needsTests :: CPattern -> Bool
366needsTests (A.A _ pattern) =
367 case pattern of
368 P.Var _ ->
369 False
371 P.Anything ->
372 False
374 P.Alias _ _ ->
375 error "aliases should never reach 'isIrrelevantTo' function"
377 P.Record _ ->
378 False
380 P.Data _ _ ->
381 True
383 P.Literal _ ->
384 True
387-- PATH PICKING HEURISTICS
389smallDefaults :: [Branch] -> Path
390smallDefaults branches =
391 let
392 allPaths =
393 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
395 weightedPaths =
396 map (\path -> (path, length (filter (isIrrelevantTo path) branches))) allPaths
397 in
398 fst (List.minimumBy (compare `on` snd) weightedPaths)
401isChoicePath :: (Path, CPattern) -> Maybe Path
402isChoicePath (path, pattern) =
403 if needsTests pattern then
404 Just path
405 else
406 Nothing
409-- smallBranchingFactor :: [Branch] -> Path