@@ -7,11 +7,14 @@ module Optimize.Names
77 , generate
88 , registerKernel
99 , registerGlobal
10+ , registerFields
11+ , registerField
1012 )
1113 where
1214
1315
1416import Data.Monoid ((<>))
17+import qualified Data.Map as Map
1518import qualified Data.Set as Set
1619import qualified Data.Text as Text
1720
@@ -30,34 +33,51 @@ newtype Tracker a =
3033 Int
3134 -> Set.Set N.Name
3235 -> Set.Set Opt.Global
33− -> (Int -> Set.Set N.Name -> Set.Set Opt.Global -> a -> r)
36+ -> Map.Map N.Name Int
37+ -> (Int -> Set.Set N.Name -> Set.Set Opt.Global -> Map.Map N.Name Int -> a -> r)
3438 -> r
3539 )
3640
3741
38−run :: Tracker a -> (Set.Set N.Name, Set.Set Opt.Global, a)
42+run :: Tracker a -> (Set.Set N.Name, Set.Set Opt.Global, Map.Map N.Name Int, a)
3943run (Tracker k) =
40− k 0 Set.empty Set.empty
41− (\_uid kernels globals value -> (kernels, globals, value))
44+ k 0 Set.empty Set.empty Map.empty
45+ (\_uid kernels globals fields value -> (kernels, globals, fields, value))
4246
4347
4448generate :: Tracker N.Name
4549generate =
46− Tracker $ \n k g ok ->
47− ok (n + 1) k g ("_n" <> Text.pack (show n))
50+ Tracker $ \n k g f ok ->
51+ ok (n + 1) k g f ("_n" <> Text.pack (show n))
4852
4953
5054registerKernel :: N.Name -> a -> Tracker a
5155registerKernel name value =
52− Tracker $ \n k g ok ->
53− ok n (Set.insert name k) g value
56+ Tracker $ \n k g f ok ->
57+ ok n (Set.insert name k) g f value
5458
5559
5660registerGlobal :: ModuleName.Canonical -> N.Name -> Tracker Opt.Expr
5761registerGlobal home name =
58− Tracker $ \n k g ok ->
62+ Tracker $ \n k g f ok ->
5963 let global = Opt.Global home name in
60− ok n k (Set.insert global g) (Opt.VarGlobal global)
64+ ok n k (Set.insert global g) f (Opt.VarGlobal global)
65+
66+
67+registerField :: N.Name -> a -> Tracker a
68+registerField name value =
69+ Tracker $ \n k g fields ok ->
70+ ok n k g (Map.insertWith (+) name 1 fields) value
71+
72+
73+registerFields :: Map.Map N.Name v -> a -> Tracker a
74+registerFields newFields value =
75+ Tracker $ \n k g fields ok ->
76+ ok n k g (Map.unionWith (+) fields (Map.map toOne newFields)) value
77+
78+
79+toOne :: a -> Int
80+toOne _ = 1
6181
6282
6383
@@ -66,40 +86,40 @@ registerGlobal home name =
6686
6787instance Functor Tracker where
6888 fmap func (Tracker kv) =
69− Tracker $ \n k g ok ->
89+ Tracker $ \n k g f ok ->
7090 let
71− ok1 n1 k1 g1 value =
72− ok n1 k1 g1 (func value)
91+ ok1 n1 k1 g1 f1 value =
92+ ok n1 k1 g1 f1 (func value)
7393 in
74− kv n k g ok1
94+ kv n k g f ok1
7595
7696
7797instance Applicative Tracker where
7898 {-# INLINE pure #-}
7999 pure value =
80− Tracker $ \n k g ok -> ok n k g value
100+ Tracker $ \n k g f ok -> ok n k g f value
81101
82102 (<*>) (Tracker kf) (Tracker kv) =
83− Tracker $ \n0 k0 g0 ok ->
103+ Tracker $ \n0 k0 g0 f0 ok ->
84104 let
85− ok1 n1 k1 g1 func =
105+ ok1 n1 k1 g1 f1 func =
86106 let
87− ok2 n2 k2 g2 value =
88− ok n2 k2 g2 (func value)
107+ ok2 n2 k2 g2 f2 value =
108+ ok n2 k2 g2 f2 (func value)
89109 in
90− kv n1 k1 g1 ok2
110+ kv n1 k1 g1 f1 ok2
91111 in
92− kf n0 k0 g0 ok1
112+ kf n0 k0 g0 f0 ok1
93113
94114
95115instance Monad Tracker where
96116 return = pure
97117
98118 (>>=) (Tracker ka) callback =
99− Tracker $ \n k g ok ->
119+ Tracker $ \n k g f ok ->
100120 let
101− ok1 n1 k1 g1 a =
121+ ok1 n1 k1 g1 f1 a =
102122 case callback a of
103− Tracker kb -> kb n1 k1 g1 ok
123+ Tracker kb -> kb n1 k1 g1 f1 ok
104124 in
105− ka n k g ok1
125+ ka n k g f ok1