{-# 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