blob: d9fe661b3c0046ba3cb9a96c84b2673f5eb0db42 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
{-# LANGUAGE GADTs #-}
module Count where
import Data.GADT.Compare
import Data.Type.Equality
import AST
data Layout = LyLeaf Integer | LyPair Layout Layout
deriving (Show)
instance Semigroup Layout where
LyLeaf a <> LyLeaf b = LyLeaf (a + b)
l@(LyLeaf _) <> LyPair l1 l2 = LyPair (l <> l1) (l <> l2)
LyPair l1 l2 <> l@(LyLeaf _) = LyPair (l1 <> l) (l2 <> l)
LyPair l1 l2 <> LyPair l3 l4 = LyPair (l1 <> l3) (l2 <> l4)
instance Monoid Layout where
mempty = LyLeaf 0
-- | Returns the maximum usage count of any component of the layout.
lycontract :: Layout -> Integer
lycontract (LyLeaf n) = n
lycontract (LyPair ly1 ly2) = max (lycontract ly1) (lycontract ly2)
-- | Given the usage count of a pair, return the usage count of its first component
lyfst :: Layout -> Layout
lyfst (LyLeaf n) = LyLeaf n
lyfst (LyPair ly _) = ly
-- | Given the usage count of a pair, return the usage count of its second component
lysnd :: Layout -> Layout
lysnd (LyLeaf n) = LyLeaf n
lysnd (LyPair _ ly) = ly
-- | Returns the usage count of an expression given its usage counts in two
-- mutually exclusive program branches.
lymax :: Layout -> Layout -> Layout
lymax (LyLeaf n) (LyLeaf m) = LyLeaf (max n m)
lymax (LyPair a1 a2) (LyPair b1 b2) = LyPair (lymax a1 b1) (lymax a2 b2)
lymax (LyLeaf n) ly@LyPair{} = lymax (LyPair (LyLeaf n) (LyLeaf n)) ly
lymax ly@LyPair{} (LyLeaf n) = lymax ly (LyPair (LyLeaf n) (LyLeaf n))
-- | Count the uses of a variable in an expression
usesOf :: Idx env t -> Exp env a -> Integer
usesOf x e = lycontract (usesOf' PathStart x e)
-- Path upwards in the tuple, starting at Start.
data Path part t where
PathFst :: Path part t -> Path part (t, a)
PathSnd :: Path part t -> Path part (a, t)
PathStart :: Path part part
-- | Count the uses of the components of a variable in an expression
usesOf' :: Path large a -> Idx env t -> Exp env a -> Layout
usesOf' _ i (App a b) = usesOf' PathStart i a <> usesOf' PathStart i b
usesOf' _ i (Lam _ e) = usesOf' PathStart (Succ i) e
usesOf' pt i (Var _ i') =
let leaf = case geq i i' of Just Refl -> LyLeaf 1
Nothing -> mempty
build :: Path a large -> Layout -> Layout
build (PathFst pt') ly = build pt' (LyPair ly mempty)
build (PathSnd pt') ly = build pt' (LyPair mempty ly)
build PathStart ly = ly
in build pt leaf
usesOf' _ i (Let a b) = usesOf' PathStart i a <> usesOf' PathStart (Succ i) b
usesOf' _ _ (Lit _) = mempty
usesOf' pt i (Cond p l r) =
usesOf' PathStart i p <> lymax (usesOf' pt i l) (usesOf' pt i r)
usesOf' _ _ (Const _) = mempty
usesOf' _ i (Pair a b) = usesOf' PathStart i a <> usesOf' PathStart i b
usesOf' pt i (Fst e) = usesOf' (PathFst pt) i e
usesOf' pt i (Snd e) = usesOf' (PathSnd pt) i e
-- TODO: Huge hack that allows arbitrary computational complexity increase.
-- The code current counts usages in the lambdas in Build and Ifold _once_,
-- whereas those usages are actually evaluated many times. The correct fix
-- would be to analyse whether the accessed indexes of the counted variable are
-- all disjoint, but that is hard.
usesOf' _ i (Build _ she fe) = usesOf' PathStart i she <> usesOf' PathStart i fe
-- usesOf' _ i (Build STZ she fe) = usesOf' PathStart i she <> usesOf' PathStart i fe
-- usesOf' _ i (Build _ she fe) = usesOf' PathStart i she <> mul2 (usesOf' PathStart i fe)
usesOf' _ i (Ifold _ fe e0 she) =
usesOf' PathStart i fe <> usesOf' PathStart i e0 <> usesOf' PathStart i she
-- usesOf' _ i (Ifold STZ fe e0 she) =
-- usesOf' PathStart i fe <> usesOf' PathStart i e0 <> usesOf' PathStart i she
-- usesOf' _ i (Ifold _ fe e0 she) =
-- mul2 (usesOf' PathStart i fe) <> usesOf' PathStart i e0 <> usesOf' PathStart i she
usesOf' _ i (Index a b) = usesOf' PathStart i a <> usesOf' PathStart i b
usesOf' _ i (Shape e) = usesOf' PathStart i e
usesOf' _ _ (Undef _) = mempty
-- mul2 :: Semigroup m => Layout m -> Layout m
-- mul2 ly = ly <> ly
|
