git.iko.soy/elk
on ebe2bb5 help sign in
elk/src/Optimize/Case.hs
173 lines 4.4 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Case (optimize) where
4import Control.Arrow (second)
5import qualified Data.Map as Map
6import Data.Map ((!))
7import qualified Data.Maybe as Maybe
9import qualified AST.Expression.Optimized as Opt
10import qualified AST.Pattern as P
11import qualified Optimize.DecisionTree as DT
15-- OPTIMIZE A CASE EXPRESSION
18optimize :: DT.VariantDict -> String -> [(P.CanonicalPattern, Opt.Expr)] -> Opt.Expr
19optimize variantDict exprName optBranches =
20 let
21 (patterns, indexedBranches) =
22 unzip (zipWith indexify [0..] optBranches)
24 decisionTree =
25 DT.compile variantDict patterns
26 in
27 treeToExpr exprName decisionTree indexedBranches
30indexify :: Int -> (a,b) -> ((a,Int), (Int,b))
31indexify index (pattern, branch) =
32 ( (pattern, index)
33 , (index, branch)
34 )
38-- CONVERT A TREE TO AN EXPRESSION
41treeToExpr :: String -> DT.DecisionTree -> [(Int, Opt.Expr)] -> Opt.Expr
42treeToExpr name decisionTree allJumps =
43 let
44 decider =
45 treeToDecider decisionTree
47 targetCounts =
48 countTargets decider
50 (choices, maybeJumps) =
51 unzip (map (createChoices targetCounts) allJumps)
52 in
53 Opt.Case
54 name
55 (insertChoices (Map.fromList choices) decider)
56 (Maybe.catMaybes maybeJumps)
60-- TREE TO DECIDER
61--
62-- Decision trees may have some redundancies, so we convert them to a Decider
63-- which has special constructs to avoid code duplication when possible.
66treeToDecider :: DT.DecisionTree -> Opt.Decider Int
67treeToDecider tree =
68 case tree of
69 DT.Match target ->
70 Opt.Leaf target
72 -- zero options
73 DT.Decision _ [] Nothing ->
74 error "compiler bug, somehow created an empty decision tree"
76 -- one option
77 DT.Decision _ [(_, subTree)] Nothing ->
78 treeToDecider subTree
80 DT.Decision _ [] (Just subTree) ->
81 treeToDecider subTree
83 -- two options
84 DT.Decision path [(test, successTree)] (Just failureTree) ->
85 toChain path test successTree failureTree
87 DT.Decision path [(test, successTree), (_, failureTree)] Nothing ->
88 toChain path test successTree failureTree
90 -- many options
91 DT.Decision path edges Nothing ->
92 let
93 (necessaryTests, fallback) =
94 (init edges, snd (last edges))
95 in
96 Opt.FanOut
97 path
98 (map (second treeToDecider) necessaryTests)
99 (treeToDecider fallback)
101 DT.Decision path edges (Just fallback) ->
102 Opt.FanOut path (map (second treeToDecider) edges) (treeToDecider fallback)
105toChain :: DT.Path -> DT.Test -> DT.DecisionTree -> DT.DecisionTree -> Opt.Decider Int
106toChain path test successTree failureTree =
107 let
108 failure =
109 treeToDecider failureTree
110 in
111 case treeToDecider successTree of
112 Opt.Chain testChain success subFailure | failure == subFailure ->
113 Opt.Chain ((path, test) : testChain) success failure
115 success ->
116 Opt.Chain [(path, test)] success failure
120-- INSERT CHOICES
121--
122-- If a target appears exactly once in a Decider, the corresponding expression
123-- can be inlined. Whether things are inlined or jumps is called a "choice".
126countTargets :: Opt.Decider Int -> Map.Map Int Int
127countTargets decisionTree =
128 case decisionTree of
129 Opt.Leaf target ->
130 Map.singleton target 1
132 Opt.Chain _ success failure ->
133 Map.unionWith (+) (countTargets success) (countTargets failure)
135 Opt.FanOut _ tests fallback ->
136 Map.unionsWith (+) (map countTargets (fallback : map snd tests))
139createChoices
140 :: Map.Map Int Int
141 -> (Int, Opt.Expr)
142 -> ( (Int, Opt.Choice), Maybe (Int, Opt.Expr) )
143createChoices targetCounts (target, branch) =
144 if targetCounts ! target == 1 then
145 ( (target, Opt.Inline branch)
146 , Nothing
147 )
149 else
150 ( (target, Opt.Jump target)
151 , Just (target, branch)
152 )
155insertChoices
156 :: Map.Map Int Opt.Choice
157 -> Opt.Decider Int
158 -> Opt.Decider Opt.Choice
159insertChoices choiceDict decider =
160 let
161 go =
162 insertChoices choiceDict
163 in
164 case decider of
165 Opt.Leaf target ->
166 Opt.Leaf (choiceDict ! target)
168 Opt.Chain testChain success failure ->
169 Opt.Chain testChain (go success) (go failure)
171 Opt.FanOut path tests fallback ->
172 Opt.FanOut path (map (second go) tests) (go fallback)