git.iko.soy/elk
on f4401ac help sign in
elk/compiler/src/Optimize/Case.hs
161 lines 4.1 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Case
3 ( optimize
4 )
5 where
8import Control.Arrow (second)
9import qualified Data.Map as Map
10import Data.Map ((!))
11import qualified Data.Maybe as Maybe
12import qualified Data.Name as Name
14import qualified AST.Canonical as Can
15import qualified AST.Optimized as Opt
16import qualified Optimize.DecisionTree as DT
20-- OPTIMIZE A CASE EXPRESSION
23optimize :: Name.Name -> Name.Name -> [(Can.Pattern, Opt.Expr)] -> Opt.Expr
24optimize temp root optBranches =
25 let
26 (patterns, indexedBranches) =
27 unzip (zipWith indexify [0..] optBranches)
29 decider = treeToDecider (DT.compile patterns)
30 targetCounts = countTargets decider
32 (choices, maybeJumps) =
33 unzip (map (createChoices targetCounts) indexedBranches)
34 in
35 Opt.Case temp root
36 (insertChoices (Map.fromList choices) decider)
37 (Maybe.catMaybes maybeJumps)
40indexify :: Int -> (a,b) -> ((a,Int), (Int,b))
41indexify index (pattern, branch) =
42 ( (pattern, index)
43 , (index, branch)
44 )
48-- TREE TO DECIDER
49--
50-- Decision trees may have some redundancies, so we convert them to a Decider
51-- which has special constructs to avoid code duplication when possible.
54treeToDecider :: DT.DecisionTree -> Opt.Decider Int
55treeToDecider tree =
56 case tree of
57 DT.Match target ->
58 Opt.Leaf target
60 -- zero options
61 DT.Decision _ [] Nothing ->
62 error "compiler bug, somehow created an empty decision tree"
64 -- one option
65 DT.Decision _ [(_, subTree)] Nothing ->
66 treeToDecider subTree
68 DT.Decision _ [] (Just subTree) ->
69 treeToDecider subTree
71 -- two options
72 DT.Decision path [(test, successTree)] (Just failureTree) ->
73 toChain path test successTree failureTree
75 DT.Decision path [(test, successTree), (_, failureTree)] Nothing ->
76 toChain path test successTree failureTree
78 -- many options
79 DT.Decision path edges Nothing ->
80 let
81 (necessaryTests, fallback) =
82 (init edges, snd (last edges))
83 in
84 Opt.FanOut
85 path
86 (map (second treeToDecider) necessaryTests)
87 (treeToDecider fallback)
89 DT.Decision path edges (Just fallback) ->
90 Opt.FanOut path (map (second treeToDecider) edges) (treeToDecider fallback)
93toChain :: DT.Path -> DT.Test -> DT.DecisionTree -> DT.DecisionTree -> Opt.Decider Int
94toChain path test successTree failureTree =
95 let
96 failure =
97 treeToDecider failureTree
98 in
99 case treeToDecider successTree of
100 Opt.Chain testChain success subFailure | failure == subFailure ->
101 Opt.Chain ((path, test) : testChain) success failure
103 success ->
104 Opt.Chain [(path, test)] success failure
108-- INSERT CHOICES
109--
110-- If a target appears exactly once in a Decider, the corresponding expression
111-- can be inlined. Whether things are inlined or jumps is called a "choice".
114countTargets :: Opt.Decider Int -> Map.Map Int Int
115countTargets decisionTree =
116 case decisionTree of
117 Opt.Leaf target ->
118 Map.singleton target 1
120 Opt.Chain _ success failure ->
121 Map.unionWith (+) (countTargets success) (countTargets failure)
123 Opt.FanOut _ tests fallback ->
124 Map.unionsWith (+) (map countTargets (fallback : map snd tests))
127createChoices
128 :: Map.Map Int Int
129 -> (Int, Opt.Expr)
130 -> ( (Int, Opt.Choice), Maybe (Int, Opt.Expr) )
131createChoices targetCounts (target, branch) =
132 if targetCounts ! target == 1 then
133 ( (target, Opt.Inline branch)
134 , Nothing
135 )
137 else
138 ( (target, Opt.Jump target)
139 , Just (target, branch)
140 )
143insertChoices
144 :: Map.Map Int Opt.Choice
145 -> Opt.Decider Int
146 -> Opt.Decider Opt.Choice
147insertChoices choiceDict decider =
148 let
149 go =
150 insertChoices choiceDict
151 in
152 case decider of
153 Opt.Leaf target ->
154 Opt.Leaf (choiceDict ! target)
156 Opt.Chain testChain success failure ->
157 Opt.Chain testChain (go success) (go failure)
159 Opt.FanOut path tests fallback ->
160 Opt.FanOut path (map (second go) tests) (go fallback)