Select0.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Data.Succinct.RankSelect.Base.Select0
    ( Select0(..)
    ) where

import Data.Succinct.RankSelect.Base.Select1
import Data.Word
import HaskellWorks.Data.AtIndex
import HaskellWorks.Data.Bits.BitShown
import HaskellWorks.Data.Bits.BitWise
import HaskellWorks.Data.Bits.ElemFixedBitSize
import HaskellWorks.Data.Bits.PopCount.PopCount0
import HaskellWorks.Data.Positioning

import qualified Data.Vector          as DV
import qualified Data.Vector.Storable as DVS

{- HLINT ignore "Reduce duplication"  -}

class Select0 v where
  -- | Find length of the shortest prefix of the given prefix that contains specified number of occurences of the bit @1@
  --
  -- If the bitstring does not have enough occurences of bit @0@ is insufficient to satisfy the query the result is undefined.
  --
  -- >>> import HaskellWorks.Data.Bits.BitRead
  -- >>> :set -XTypeApplications
  --
  -- >>> select0 (unsafeBitRead @Word8 "11111111") 0
  -- 0
  -- >>> select0 (unsafeBitRead @Word8 "11110111") 1
  -- 5
  -- >>> select0 (unsafeBitRead @Word8 "00000000") 4
  -- 4
  -- >>> select0 (unsafeBitRead @Word8 "11000000") 4
  -- 6
  -- >>> select0 (unsafeBitRead @Word8 "01100000") 4
  -- 6
  select0
    :: v      -- ^ The bitstring
    -> Count  -- ^ The number of zeros
    -> Count

deriving instance Select0 a => Select0 (BitShown a)

-- TODO: Implement NOT in terms of select for word-16
instance Select0 Word8 where
  select0 v = select1 (comp v)
  {-# INLINABLE select0 #-}

instance Select0 Word16 where
  select0 v = select1 (comp v)
  {-# INLINABLE select0 #-}

instance Select0 Word32 where
  select0 v = select1 (comp v)
  {-# INLINABLE select0 #-}

instance Select0 Word64 where
  select0 v = select1 (comp v)
  {-# INLINABLE select0 #-}

instance Select0 [Bool] where
  select0 = go 0
    where go r _ 0          = r
          go r (False:bs) c = go (r + 1) bs (c - 1)
          go r (True:bs)  c = go (r + 1) bs  c
          go _ []         _ = error "Out of range"
  {-# INLINABLE select0 #-}

instance Select0 [Word8] where
  select0 v c = go v c 0
    where go :: [Word8] -> Count -> Count -> Count
          go _ 0  acc = acc
          go u d acc = let w = head u in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (tail u) (d - pc) (acc + elemFixedBitSize u)
  {-# INLINABLE select0 #-}

instance Select0 [Word16] where
  select0 v c = go v c 0
    where go :: [Word16] -> Count -> Count -> Count
          go _ 0  acc = acc
          go u d acc = let w = head u in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (tail u) (d - pc) (acc + elemFixedBitSize u)
  {-# INLINABLE select0 #-}

instance Select0 [Word32] where
  select0 v c = go v c 0
    where go :: [Word32] -> Count -> Count -> Count
          go _ 0  acc = acc
          go u d acc = let w = head u in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (tail u) (d - pc) (acc + elemFixedBitSize u)
  {-# INLINABLE select0 #-}

instance Select0 [Word64] where
  select0 v c = go v c 0
    where go :: [Word64] -> Count -> Count -> Count
          go _ 0  acc = acc
          go u d acc = let w = head u in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (tail u) (d - pc) (acc + elemFixedBitSize u)
  {-# INLINABLE select0 #-}

instance Select0 (DV.Vector Word8) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DV.Vector Word16) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DV.Vector Word32) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DV.Vector Word64) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DVS.Vector Word8) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DVS.Vector Word16) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DVS.Vector Word32) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}

instance Select0 (DVS.Vector Word64) where
  select0 v c = go 0 c 0
    where go _ 0  acc = acc
          go n d acc = let w = (v !!! n) in
            case popCount0 w of
              pc | d <= pc -> select0 w d + acc
              pc           -> go (n + 1) (d - pc) (acc + elemFixedBitSize v)
  {-# INLINABLE select0 #-}