ThingsHEADmaster

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