diff options
Diffstat (limited to 'Count.hs')
| -rw-r--r-- | Count.hs | 94 |
1 files changed, 94 insertions, 0 deletions
diff --git a/Count.hs b/Count.hs new file mode 100644 index 0000000..d9fe661 --- /dev/null +++ b/Count.hs @@ -0,0 +1,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 |