git.iko.soy/elk
on dbf97e2 help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
417 lines 10 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 ->
319 if test == Constructor name then
320 Just (Branch goal (start ++ subPositions path patterns ++ end))
322 else
323 Nothing
325 P.Literal lit ->
326 if test == Literal lit then
327 Just (Branch goal (start ++ end))
329 else
330 Nothing
332 _ ->
333 Just branch
335 _ ->
336 Just branch
339extract
340 :: Path
341 -> [(Path, CPattern)]
342 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
343extract selectedPath pathPatterns =
344 case pathPatterns of
345 [] ->
346 Nothing
348 first@(path, pattern) : rest ->
349 if path == selectedPath then
350 Just ([], pattern, rest)
352 else
353 case extract selectedPath rest of
354 Nothing ->
355 Nothing
357 Just (start, foundPattern, end) ->
358 Just (first : start, foundPattern, end)
361-- FIND IRRELEVANT BRANCHES
363isIrrelevantTo :: Path -> Branch -> Bool
364isIrrelevantTo selectedPath (Branch _ pathPatterns) =
365 case List.lookup selectedPath pathPatterns of
366 Nothing ->
367 True
369 Just pattern ->
370 not (needsTests pattern)
373needsTests :: CPattern -> Bool
374needsTests (A.A _ pattern) =
375 case pattern of
376 P.Var _ ->
377 False
379 P.Anything ->
380 False
382 P.Alias _ _ ->
383 error "aliases should never reach 'isIrrelevantTo' function"
385 P.Record _ ->
386 False
388 P.Data _ _ ->
389 True
391 P.Literal _ ->
392 True
395-- PATH PICKING HEURISTICS
397smallDefaults :: [Branch] -> Path
398smallDefaults branches =
399 let
400 allPaths =
401 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
403 weightedPaths =
404 map (\path -> (path, length (filter (isIrrelevantTo path) branches))) allPaths
405 in
406 fst (List.minimumBy (compare `on` snd) weightedPaths)
409isChoicePath :: (Path, CPattern) -> Maybe Path
410isChoicePath (path, pattern) =
411 if needsTests pattern then
412 Just path
413 else
414 Nothing
417-- smallBranchingFactor :: [Branch] -> Path