{-# OPTIONS_HADDOCK hide #-}

module PopKey.Internal1 where

import Control.Monad.ST
import Data.Bit as B
import qualified Data.ByteString as BS
import Data.Foldable
import Data.STRef
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as MUV
import GHC.Generics (Generic)
import GHC.Word
import HaskellWorks.Data.Bits.PopCount.PopCount1
import qualified HaskellWorks.Data.RankSelect.Base.Select1
import HaskellWorks.Data.RankSelect.CsPoppy
import Unsafe.Coerce


data PKPrim =
    ConstSize !BS.ByteString {-# UNPACK #-} !Word32 {-# UNPACK #-} !Word32 -- raw / item size / item count
  | Var !CsPoppy !BS.ByteString {-# UNPACK #-} !Word32 {-# UNPACK #-} !Word32 -- poppy / raw / min size / step size
  deriving (forall x. Rep PKPrim x -> PKPrim
forall x. PKPrim -> Rep PKPrim x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PKPrim x -> PKPrim
$cfrom :: forall x. PKPrim -> Rep PKPrim x
Generic,PKPrim -> PKPrim -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PKPrim -> PKPrim -> Bool
$c/= :: PKPrim -> PKPrim -> Bool
== :: PKPrim -> PKPrim -> Bool
$c== :: PKPrim -> PKPrim -> Bool
Eq)


-- 0-based indexing, for my sanity
{-# INLINE select1' #-}
select1' :: CsPoppy -> Int -> Int
select1' :: CsPoppy -> Int -> Int
select1' CsPoppy
p Int
i =
  forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall v. Select1 v => v -> Count -> Count
HaskellWorks.Data.RankSelect.Base.Select1.select1 CsPoppy
p (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i forall a. Num a => a -> a -> a
+ Count
1)) forall a. Num a => a -> a -> a
- Int
1

{-# INLINABLE pkLength #-}
pkLength :: PKPrim -> Int
pkLength :: PKPrim -> Int
pkLength (ConstSize ByteString
_ Word32
_ Word32
l) = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
l
pkLength (Var CsPoppy
p ByteString
_ Word32
_ Word32
_) = (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\Count
x -> Count
x forall a. Num a => a -> a -> a
- Count
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. PopCount1 v => v -> Count
popCount1) CsPoppy
p

{-# INLINABLE pkIndex #-}
pkIndex :: PKPrim -> Int -> BS.ByteString
pkIndex :: PKPrim -> Int -> ByteString
pkIndex (ConstSize ByteString
r (forall a b. (Integral a, Num b) => a -> b
fromIntegral -> Int
s) Word32
_) Int
i = if Int
s forall a. Eq a => a -> a -> Bool
== Int
0 then forall a. Monoid a => a
mempty else Int -> ByteString -> ByteString
BS.take Int
s (Int -> ByteString -> ByteString
BS.drop (Int
i forall a. Num a => a -> a -> a
* Int
s) ByteString
r)
pkIndex (Var CsPoppy
p ByteString
r (forall a b. (Integral a, Num b) => a -> b
fromIntegral -> Int
minSize) (forall a b. (Integral a, Num b) => a -> b
fromIntegral -> Int
step)) Int
i = do
  let Int
o :: Int = CsPoppy -> Int -> Int
select1' CsPoppy
p Int
i
      Int
d :: Int = CsPoppy -> Int -> Int
select1' CsPoppy
p (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a. Num a => a -> a -> a
- Int
o

  Int -> ByteString -> ByteString
BS.take (Int
minSize forall a. Num a => a -> a -> a
+ Int
step forall a. Num a => a -> a -> a
* (Int
d forall a. Num a => a -> a -> a
- Int
1)) (Int -> ByteString -> ByteString
BS.drop (Int
step forall a. Num a => a -> a -> a
* (Int
o forall a. Num a => a -> a -> a
- Int
i) forall a. Num a => a -> a -> a
+ Int
i forall a. Num a => a -> a -> a
* Int
minSize) ByteString
r)

makePK :: [ BS.ByteString ] -> PKPrim
makePK :: [ByteString] -> PKPrim
makePK [] = ByteString -> Word32 -> Word32 -> PKPrim
ConstSize forall a. Monoid a => a
mempty Word32
0 Word32
0
makePK [ByteString]
bs = forall a. (forall s. ST s a) -> a
runST do
  let minSize :: Int
minSize = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (ByteString -> Int
BS.length forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ByteString]
bs)
      step :: Int
step = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
a ByteString
x -> forall a. Integral a => a -> a -> a
gcd (ByteString -> Int
BS.length ByteString
x forall a. Num a => a -> a -> a
- Int
minSize) Int
a) (ByteString -> Int
BS.length (forall a. [a] -> a
head [ByteString]
bs) forall a. Num a => a -> a -> a
- Int
minSize) [ByteString]
bs

  if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int
minSizeforall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Int
BS.length) [ByteString]
bs
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure do ByteString -> Word32 -> Word32 -> PKPrim
ConstSize ([ByteString] -> ByteString
BS.concat [ByteString]
bs) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
minSize) (forall a b. (Integral a, Num b) => a -> b
fromIntegral do forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
bs)
     else do
       -- raw indexing vector
       Vector Bit
bv :: UV.Vector Bit <- do
         MVector s Bit
v <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new do Int
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
a ByteString
x -> Int
a forall a. Num a => a -> a -> a
+ Int
1 forall a. Num a => a -> a -> a
+ (ByteString -> Int
BS.length ByteString
x forall a. Num a => a -> a -> a
- Int
minSize) forall a. Integral a => a -> a -> a
`div` Int
step) Int
0 [ByteString]
bs
         forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.unsafeWrite MVector s Bit
v Int
0 Bit
1

         STRef s Int
base_ref <- forall a s. a -> ST s (STRef s a)
newSTRef Int
0

         forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ByteString]
bs \ByteString
x -> do
           let d :: Int
d = ((ByteString -> Int
BS.length ByteString
x forall a. Num a => a -> a -> a
- Int
minSize) forall a. Integral a => a -> a -> a
`div` Int
step) forall a. Num a => a -> a -> a
+ Int
1
           Int
b <- forall s a. STRef s a -> ST s a
readSTRef STRef s Int
base_ref
           forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.unsafeWrite MVector s Bit
v (Int
b forall a. Num a => a -> a -> a
+ Int
d) Bit
1
           forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
base_ref (Int
b forall a. Num a => a -> a -> a
+ Int
d)

         forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
UV.unsafeFreeze MVector s Bit
v

       let Vector Count
uv64 :: UV.Vector Word64 = forall a b. a -> b
unsafeCoerce do Vector Bit -> Vector Word
cloneToWords Vector Bit
bv
           Vector Count
sv64 :: SV.Vector Word64 = forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
SV.convert Vector Count
uv64

           CsPoppy
ppy :: CsPoppy = Vector Count -> CsPoppy
makeCsPoppy Vector Count
sv64

       forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ CsPoppy -> ByteString -> Word32 -> Word32 -> PKPrim
Var CsPoppy
ppy ([ByteString] -> ByteString
BS.concat [ByteString]
bs) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
minSize) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
step)

-- returns @-1@ if not found
{-# INLINABLE bin_search #-}
bin_search :: PKPrim -> BS.ByteString -> Int -> Int -> Int
bin_search :: PKPrim -> ByteString -> Int -> Int -> Int
bin_search PKPrim
vs ByteString
q = Int -> Int -> Int
go
  where
    go :: Int -> Int -> Int
    go :: Int -> Int -> Int
go Int
l Int
r
      | Int
r forall a. Ord a => a -> a -> Bool
>= Int
l = do
          let m :: Int
m = Int
l forall a. Num a => a -> a -> a
+ (Int
r forall a. Num a => a -> a -> a
- Int
l) forall a. Integral a => a -> a -> a
`div` Int
2
              p :: ByteString
p = PKPrim -> Int -> ByteString
pkIndex PKPrim
vs Int
m
          if ByteString
p forall a. Ord a => a -> a -> Bool
> ByteString
q
             then Int -> Int -> Int
go Int
l (Int
m forall a. Num a => a -> a -> a
- Int
1)
             else if ByteString
p forall a. Eq a => a -> a -> Bool
== ByteString
q
                     then Int
m
                     else Int -> Int -> Int
go (Int
m forall a. Num a => a -> a -> a
+ Int
1) Int
r
      | Bool
otherwise = -Int
1