{-# OPTIONS_HADDOCK hide #-}

module PopKey.Internal2 where

import Control.Monad.ST
import Data.Bit as B
import qualified Data.ByteString as BS
import Data.Either
import Data.Foldable
import HaskellWorks.Data.RankSelect.CsPoppy
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as MUV
import GHC.Word
import HaskellWorks.Data.Bits.BitWise ((.?.))
import HaskellWorks.Data.RankSelect.Base.Rank0
import HaskellWorks.Data.RankSelect.Base.Rank1
import Unsafe.Coerce

import PopKey.Internal1


data F s a where
  Single :: !a -> F () a
  Prod :: !(F s1 a) -> !(F s2 a) -> F (s1 , s2) a
  Sum :: {-# UNPACK #-} !Word32 -> CsPoppy -> !(F s1 a) -> !(F s2 a) -> F (Either s1 s2) a
  -- cardinality / poppy ; poppy undefined if cardinality = 0
  -- 0 indicates storage in the left / 1 indicates storage in the right

data F' s a where
  Single' :: a -> F' () a
  Prod' :: (F' s1 a) -> (F' s2 a) -> F' (s1 , s2) a
  Sum' :: (Either (F' s1 a) (F' s2 a)) -> F' (Either s1 s2) a

instance Eq a => Eq (F' s a) where
  {-# INLINEABLE (==) #-}
  == :: F' s a -> F' s a -> Bool
(==) (Single' a
x) (Single' a
y) = a
x forall a. Eq a => a -> a -> Bool
== a
y
  (==) (Prod' F' s1 a
x1 F' s2 a
y1) (Prod' F' s1 a
x2 F' s2 a
y2) = (F' s1 a
x1 forall a. Eq a => a -> a -> Bool
== F' s1 a
x2) Bool -> Bool -> Bool
&& (F' s2 a
y1 forall a. Eq a => a -> a -> Bool
== F' s2 a
y2)
  (==) (Sum' Either (F' s1 a) (F' s2 a)
s1) (Sum' Either (F' s1 a) (F' s2 a)
s2) = Either (F' s1 a) (F' s2 a)
s1 forall a. Eq a => a -> a -> Bool
== Either (F' s1 a) (F' s2 a)
s2

instance Ord a => Ord (F' s a) where
  {-# INLINABLE (<=) #-}
  <= :: F' s a -> F' s a -> Bool
(<=) (Single' a
x) (Single' a
y) = a
x forall a. Ord a => a -> a -> Bool
<= a
y
  (<=) (Prod' F' s1 a
x1 F' s2 a
y1) (Prod' F' s1 a
x2 F' s2 a
y2) = (F' s1 a
x1 , F' s2 a
y1) forall a. Ord a => a -> a -> Bool
<= (F' s1 a
x2 , F' s2 a
y2)
  (<=) (Sum' Either (F' s1 a) (F' s2 a)
s1) (Sum' Either (F' s1 a) (F' s2 a)
s2) = Either (F' s1 a) (F' s2 a)
s1 forall a. Ord a => a -> a -> Bool
<= Either (F' s1 a) (F' s2 a)
s2

flength :: F s PKPrim -> Int
flength :: forall s. F s PKPrim -> Int
flength (Single PKPrim
a) = PKPrim -> Int
pkLength PKPrim
a
flength (Prod F s1 PKPrim
x F s2 PKPrim
_) = forall s. F s PKPrim -> Int
flength F s1 PKPrim
x
flength (Sum Word32
l CsPoppy
_ F s1 PKPrim
_ F s2 PKPrim
_) = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
l

data I s where
  ISingle :: I ()
  IProd :: I s1 -> I s2 -> I (s1 , s2)
  ISum :: I s1 -> I s2 -> I (Either s1 s2)

-- index must be valid
rawq :: Int -> F s PKPrim -> F' s BS.ByteString
rawq :: forall s. Int -> F s PKPrim -> F' s ByteString
rawq Int
i = forall s. F s PKPrim -> F' s ByteString
go
  where
    go :: F s PKPrim -> F' s BS.ByteString
    go :: forall s. F s PKPrim -> F' s ByteString
go (Single PKPrim
pk) = forall a. a -> F' () a
Single' (PKPrim -> Int -> ByteString
pkIndex PKPrim
pk Int
i)
    go (Prod F s1 PKPrim
x F s2 PKPrim
y) = forall s1 a s2. F' s1 a -> F' s2 a -> F' (s1, s2) a
Prod' (forall s. F s PKPrim -> F' s ByteString
go F s1 PKPrim
x) (forall s. F s PKPrim -> F' s ByteString
go F s2 PKPrim
y)
    go (Sum Word32
_ CsPoppy
pk F s1 PKPrim
l F s2 PKPrim
r) = do
      let b1pos :: Position
b1pos = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
      if CsPoppy
pk forall a. TestBit a => a -> Position -> Bool
.?. Position
b1pos -- nasty! this uses 0-based indexing, while rank/select use 1-based indexing
         then forall s1 a s2. Either (F' s1 a) (F' s2 a) -> F' (Either s1 s2) a
Sum' (forall a b. b -> Either a b
Right (forall s. Int -> F s PKPrim -> F' s ByteString
rawq (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall v. Rank1 v => v -> Count -> Count
rank1 CsPoppy
pk (forall a b. (Integral a, Num b) => a -> b
fromIntegral Position
b1pos))) F s2 PKPrim
r))
         else forall s1 a s2. Either (F' s1 a) (F' s2 a) -> F' (Either s1 s2) a
Sum' (forall a b. a -> Either a b
Left (forall s. Int -> F s PKPrim -> F' s ByteString
rawq (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall v. Rank0 v => v -> Count -> Count
rank0 CsPoppy
pk (forall a b. (Integral a, Num b) => a -> b
fromIntegral Position
b1pos))) F s1 PKPrim
l))

-- returns @-1@ if not found
{-# INLINABLE bin_search2 #-}
bin_search2 :: F s PKPrim -> F' s BS.ByteString -> Int -> Int -> Int
bin_search2 :: forall s. F s PKPrim -> F' s ByteString -> Int -> Int -> Int
bin_search2 F s PKPrim
vs F' s 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 :: F' s ByteString
p = forall s. Int -> F s PKPrim -> F' s ByteString
rawq Int
m F s PKPrim
vs
          if F' s ByteString
p forall a. Ord a => a -> a -> Bool
> F' s ByteString
q
             then Int -> Int -> Int
go Int
l (Int
m forall a. Num a => a -> a -> a
- Int
1)
             else if F' s ByteString
p forall a. Eq a => a -> a -> Bool
== F' s 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

{-# INLINE query #-}
query :: forall a s . (F' s BS.ByteString -> a) -> F s PKPrim -> Int -> a
query :: forall a s. (F' s ByteString -> a) -> F s PKPrim -> Int -> a
query F' s ByteString -> a
d F s PKPrim
pk Int
i = F' s ByteString -> a
d (forall s. Int -> F s PKPrim -> F' s ByteString
rawq Int
i F s PKPrim
pk)

{-# INLINABLE construct #-}
construct :: forall a s f . Foldable f
          => I s
          -> (a -> F' s BS.ByteString) 
          -> f a
          -> F s PKPrim
construct :: forall a s (f :: * -> *).
Foldable f =>
I s -> (a -> F' s ByteString) -> f a -> F s PKPrim
construct = \I s
s a -> F' s ByteString
e f a
f -> if forall (t :: * -> *) a. Foldable t => t a -> Int
length f a
f forall a. Eq a => a -> a -> Bool
== Int
0
  then forall t. I t -> F t PKPrim
fancyZero I s
s
  else forall t. I t -> [F' t ByteString] -> F t PKPrim
go I s
s (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> F' s ByteString
e) forall a. Monoid a => a
mempty f a
f)
  where
    fancyZero :: forall t . I t -> F t PKPrim
    fancyZero :: forall t. I t -> F t PKPrim
fancyZero I t
ISingle = forall a. a -> F () a
Single (ByteString -> Word32 -> Word32 -> PKPrim
ConstSize forall a. Monoid a => a
mempty Word32
0 Word32
0)
    fancyZero (IProd I s1
x I s2
y) = forall s1 a s2. F s1 a -> F s2 a -> F (s1, s2) a
Prod (forall t. I t -> F t PKPrim
fancyZero I s1
x) (forall t. I t -> F t PKPrim
fancyZero I s2
y)
    fancyZero (ISum I s1
x I s2
y) = forall s1 a s2.
Word32 -> CsPoppy -> F s1 a -> F s2 a -> F (Either s1 s2) a
Sum Word32
0 forall a. HasCallStack => a
undefined (forall t. I t -> F t PKPrim
fancyZero I s1
x) (forall t. I t -> F t PKPrim
fancyZero I s2
y)

    go :: forall t . I t -> [ F' t BS.ByteString ] -> F t PKPrim
    go :: forall t. I t -> [F' t ByteString] -> F t PKPrim
go I t
ISingle = \[F' t ByteString]
ys -> forall a. a -> F () a
Single ([ByteString] -> PKPrim
makePK (F' () ByteString -> ByteString
fromSingle forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [F' t ByteString]
ys))
      where
        fromSingle :: F' () BS.ByteString -> BS.ByteString
        fromSingle :: F' () ByteString -> ByteString
fromSingle (Single' ByteString
x) = ByteString
x
    go (IProd I s1
s1 I s2
s2) = \[F' t ByteString]
ys -> do
      let ([F' s1 ByteString]
as , [F' s2 ByteString]
bs) = forall a b. [(a, b)] -> ([a], [b])
unzip (forall s1 s2.
F' (s1, s2) ByteString -> (F' s1 ByteString, F' s2 ByteString)
fromProd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [F' t ByteString]
ys)
      forall s1 a s2. F s1 a -> F s2 a -> F (s1, s2) a
Prod (forall t. I t -> [F' t ByteString] -> F t PKPrim
go I s1
s1 [F' s1 ByteString]
as) (forall t. I t -> [F' t ByteString] -> F t PKPrim
go I s2
s2 [F' s2 ByteString]
bs)
      where
        fromProd :: forall s1 s2 . F' (s1 , s2) BS.ByteString
                 -> (F' s1 BS.ByteString , F' s2 BS.ByteString)
        fromProd :: forall s1 s2.
F' (s1, s2) ByteString -> (F' s1 ByteString, F' s2 ByteString)
fromProd (Prod' F' s1 ByteString
x F' s2 ByteString
y) = (F' s1 ByteString
x , F' s2 ByteString
y)
    go (ISum I s1
s1 I s2
s2) = \[F' t ByteString]
ys -> do
      let zs :: [Either (F' s1 ByteString) (F' s2 ByteString)]
zs = forall s1 s2.
F' (Either s1 s2) ByteString
-> Either (F' s1 ByteString) (F' s2 ByteString)
fromSum forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [F' t ByteString]
ys
          l :: Int
l = forall (t :: * -> *) a. Foldable t => t a -> Int
length [F' t ByteString]
ys

          Vector Bit
bv :: UV.Vector Bit = forall a. (forall s. ST s a) -> a
runST do
            MVector s Bit
v <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new Int
l

            forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (forall a b. [a] -> [b] -> [(a, b)]
zip [ Int
0 .. ] [Either (F' s1 ByteString) (F' s2 ByteString)]
zs) \(Int
i,Either (F' s1 ByteString) (F' s2 ByteString)
x) -> case Either (F' s1 ByteString) (F' s2 ByteString)
x of
              Left F' s1 ByteString
_ -> forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.unsafeWrite MVector s Bit
v Int
i Bit
0
              Right F' s2 ByteString
_ -> forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.unsafeWrite MVector s Bit
v Int
i Bit
1

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

          Vector Count
uv64 :: UV.Vector Word64 = forall a b. a -> b
unsafeCoerce do Vector Bit -> Vector Word
cloneToWords forall a b. (a -> b) -> a -> b
$ 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

          ([F' s1 ByteString]
as , [F' s2 ByteString]
bs) = forall a b. [Either a b] -> ([a], [b])
partitionEithers [Either (F' s1 ByteString) (F' s2 ByteString)]
zs
      
      forall s1 a s2.
Word32 -> CsPoppy -> F s1 a -> F s2 a -> F (Either s1 s2) a
Sum (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l) CsPoppy
ppy (forall t. I t -> [F' t ByteString] -> F t PKPrim
f I s1
s1 [F' s1 ByteString]
as) (forall t. I t -> [F' t ByteString] -> F t PKPrim
f I s2
s2 [F' s2 ByteString]
bs)
      where
        f :: I t -> [F' t ByteString] -> F t PKPrim
f I t
s [] = forall t. I t -> F t PKPrim
fancyZero I t
s
        f I t
s [F' t ByteString]
xs = forall t. I t -> [F' t ByteString] -> F t PKPrim
go I t
s [F' t ByteString]
xs
        
        fromSum :: forall s1 s2 . F' (Either s1 s2) BS.ByteString
                -> Either (F' s1 BS.ByteString) (F' s2 BS.ByteString)
        fromSum :: forall s1 s2.
F' (Either s1 s2) ByteString
-> Either (F' s1 ByteString) (F' s2 ByteString)
fromSum (Sum' Either (F' s1 ByteString) (F' s2 ByteString)
x) = Either (F' s1 ByteString) (F' s2 ByteString)
x