PopCount2.hs

{-# LANGUAGE GADTs, KindSignatures, ScopedTypeVariables, Rank2Types,
             UndecidableInstances, MultiParamTypeClasses #-}
module PopCount2 where

import CLaSH.Prelude
import CLaSH.Sized.Internal.Index
import Data.Proxy (Proxy (..))
import Data.Singletons.Prelude
import qualified Data.List as L

data Tree :: Nat -> * -> * where
  Tip    :: a -> Tree 0 a
  Branch :: Tree n a -> Tree n a -> Tree (n+1) a

-- | Convert from a vector of length @2^n@ to a perfect binary tree of depth @n@
fromVector :: KnownNat n => Vec (2 ^ n) a -> Tree n a
fromVector v = fromVector' (toUNat snat) v

fromVector' :: UNat n -> Vec (2^n) a -> Tree n a
fromVector' _         (Cons a Nil) = Tip a
fromVector' (USucc s) vs           = Branch (fromVector' s l) (fromVector' s r)
  where
    (l,r) = splitAtU (powUNat u2 s) vs
    u2    = USucc (USucc UZero)

-- | SplitAtU function that is normally hidden in "CLaSH.Sized.Vector"
splitAtU :: UNat m -> Vec (m + n) a -> (Vec m a, Vec n a)
splitAtU UZero     ys            = (Nil,ys)
splitAtU (USucc s) (y `Cons` ys) = let (as,bs) = splitAtU s ys
                                   in  (y `Cons` as, bs)

-- | A dependently typed tree-fold on perfect binary trees
tfold :: forall p k a .
         Proxy (p :: TyFun Nat * -> *)
      -> (a -> (p $ 0))
      -> (forall l . Proxy l -> (p $ l) -> (p $ l) -> (p $ (l+1)))
      -> Tree k a
      -> (p $ k)
tfold _ f _ (Tip a)                    = f a
tfold m f g (Branch (l :: Tree z a) r) = g (Proxy :: Proxy z)
                                           (tfold m f g l)
                                           (tfold m f g r)

-- | Missing @ExtendingNum@ instance for @Index@
instance ExtendingNum (Index n) (Index m) where
  type AResult (Index n) (Index m) = Index (n+m-1)
  plus (I a) (I b) = I (a + b)

-- | Doubling @Index@ motive
data IIndex (f :: TyFun Nat *) :: *
type instance Apply IIndex l = Index ((2^l)+1)

-- | Count the number of bits given a perfect tree of @Bit@s
popCountT :: KnownNat k => Tree k Bit -> Index ((2^k)+1)
popCountT = tfold (Proxy :: Proxy IIndex) fromIntegral (\_ x y -> plus x y)

-- | Count the number of bits in a BitVector of length @2^k@
popCountBV :: (KnownNat k, KnownNat (2^k)) => BitVector (2^k) -> Index ((2^k)+1)
popCountBV = popCountT . fromVector . bv2v