{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Test.Ouroboros.Consensus.ChainGenerator.BitVector (
    -- * Finding
    MaybeFound (JustFound, NothingFound)
  , findIthActiveInV
  , findIthEmptyInMV
  , findIthEmptyInV
    -- * Counting
  , countActivesInMV
  , countActivesInV
    -- * Slots
  , setMV
  , testMV
  , testV
    -- * Generating
  , SomeDensityWindow (SomeDensityWindow)
  , fillInWindow
  ) where

import           Control.Monad.ST (ST, runST)
import           Data.Functor ((<&>))
import qualified Data.Vector.Unboxed.Mutable as MV
import qualified System.Random.Stateful as R
import qualified Test.Ouroboros.Consensus.ChainGenerator.Counting as C
import qualified Test.Ouroboros.Consensus.ChainGenerator.Slot as S
import           Test.Ouroboros.Consensus.ChainGenerator.Slot
                     (E (ActiveSlotE, EmptySlotE, SlotE), POL, PreImage, S)
import qualified Test.Ouroboros.Consensus.ChainGenerator.Some as Some

-----

data MaybeFound base =
    NothingFound
  |
    JustFound
      {-# UNPACK #-} !(C.Index base SlotE)
  deriving (MaybeFound base -> MaybeFound base -> Bool
(MaybeFound base -> MaybeFound base -> Bool)
-> (MaybeFound base -> MaybeFound base -> Bool)
-> Eq (MaybeFound base)
forall base. MaybeFound base -> MaybeFound base -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall base. MaybeFound base -> MaybeFound base -> Bool
== :: MaybeFound base -> MaybeFound base -> Bool
$c/= :: forall base. MaybeFound base -> MaybeFound base -> Bool
/= :: MaybeFound base -> MaybeFound base -> Bool
Eq, ReadPrec [MaybeFound base]
ReadPrec (MaybeFound base)
Int -> ReadS (MaybeFound base)
ReadS [MaybeFound base]
(Int -> ReadS (MaybeFound base))
-> ReadS [MaybeFound base]
-> ReadPrec (MaybeFound base)
-> ReadPrec [MaybeFound base]
-> Read (MaybeFound base)
forall base. ReadPrec [MaybeFound base]
forall base. ReadPrec (MaybeFound base)
forall base. Int -> ReadS (MaybeFound base)
forall base. ReadS [MaybeFound base]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall base. Int -> ReadS (MaybeFound base)
readsPrec :: Int -> ReadS (MaybeFound base)
$creadList :: forall base. ReadS [MaybeFound base]
readList :: ReadS [MaybeFound base]
$creadPrec :: forall base. ReadPrec (MaybeFound base)
readPrec :: ReadPrec (MaybeFound base)
$creadListPrec :: forall base. ReadPrec [MaybeFound base]
readListPrec :: ReadPrec [MaybeFound base]
Read, Int -> MaybeFound base -> ShowS
[MaybeFound base] -> ShowS
MaybeFound base -> String
(Int -> MaybeFound base -> ShowS)
-> (MaybeFound base -> String)
-> ([MaybeFound base] -> ShowS)
-> Show (MaybeFound base)
forall base. Int -> MaybeFound base -> ShowS
forall base. [MaybeFound base] -> ShowS
forall base. MaybeFound base -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall base. Int -> MaybeFound base -> ShowS
showsPrec :: Int -> MaybeFound base -> ShowS
$cshow :: forall base. MaybeFound base -> String
show :: MaybeFound base -> String
$cshowList :: forall base. [MaybeFound base] -> ShowS
showList :: [MaybeFound base] -> ShowS
Show)

-- | Trivial wrapper around 'findIthEmptyInMV'
findIthEmptyInV ::
     POL   pol
  => proxy pol
  -> C.Vector base SlotE S
  -> C.Index base (PreImage pol EmptySlotE)
  -> MaybeFound base
findIthEmptyInV :: forall (pol :: Pol) (proxy :: Pol -> *) base.
POL pol =>
proxy pol
-> Vector base 'SlotE S
-> Index base (PreImage pol 'EmptySlotE)
-> MaybeFound base
findIthEmptyInV proxy pol
pol Vector base 'SlotE S
v Index base (PreImage pol 'EmptySlotE)
i =
    (forall s. ST s (MaybeFound base)) -> MaybeFound base
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (MaybeFound base)) -> MaybeFound base)
-> (forall s. ST s (MaybeFound base)) -> MaybeFound base
forall a b. (a -> b) -> a -> b
$ Vector base 'SlotE S -> ST s (MVector base 'SlotE s S)
forall {k1} {k2} a (base :: k1) (elem :: k2) s.
Unbox a =>
Vector base elem a -> ST s (MVector base elem s a)
C.unsafeThawV Vector base 'SlotE S
v ST s (MVector base 'SlotE s S)
-> (MVector base 'SlotE s S -> ST s (MaybeFound base))
-> ST s (MaybeFound base)
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MVector base 'SlotE s S
mv -> proxy pol
-> MVector base 'SlotE s S
-> Index base (PreImage pol 'EmptySlotE)
-> ST s (MaybeFound base)
forall (proxy :: Pol -> *) (pol :: Pol) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> Index base (PreImage pol 'EmptySlotE)
-> ST s (MaybeFound base)
findIthEmptyInMV proxy pol
pol MVector base 'SlotE s S
mv Index base (PreImage pol 'EmptySlotE)
i

-- | Find the (i+1)st empty slot in a window
--
-- * @findIthEmptyInMV notInverted v 0@ yields the first empty slot
-- * @findIthEmptyInMV notInverted v 1@ yields the second empty slot
-- * @findIthEmptyInMV notInverted v k@ yields the @k+1@st empty slot
--
-- > findIthEmptyInMV notInverted 01101 0 == JustFound 0
-- > findIthEmptyInMV notInverted 01101 1 == JustFound 3
-- > findIthEmptyInMV notInverted 01101 2 == NothingFound
--
findIthEmptyInMV ::
  forall proxy pol base s.
     POL   pol
  => proxy pol
  -> C.MVector base SlotE s S
  -> C.Index base (PreImage pol EmptySlotE)
  -> ST s (MaybeFound base)
findIthEmptyInMV :: forall (proxy :: Pol -> *) (pol :: Pol) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> Index base (PreImage pol 'EmptySlotE)
-> ST s (MaybeFound base)
findIthEmptyInMV proxy pol
pol MVector base 'SlotE s S
mv Index base (PreImage pol 'EmptySlotE)
i =
    Int
-> Var base (PreImage pol 'EmptySlotE) -> ST s (MaybeFound base)
go Int
0 (Index base (PreImage pol 'EmptySlotE)
-> Var base (PreImage pol 'EmptySlotE)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Var base elem
C.toVar Index base (PreImage pol 'EmptySlotE)
i)
  where
    go :: Int
-> Var base (PreImage pol 'EmptySlotE) -> ST s (MaybeFound base)
go !Int
j !Var base (PreImage pol 'EmptySlotE)
toSkip = if Count base 'SlotE Total -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount (MVector base 'SlotE s S -> Count base 'SlotE Total
forall {kelem} a base (elem :: kelem) s.
Unbox a =>
MVector base elem s a -> Size base elem
C.lengthMV MVector base 'SlotE s S
mv) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
j then MaybeFound base -> ST s (MaybeFound base)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MaybeFound base
forall base. MaybeFound base
NothingFound else do
        S
w <- MVector base 'SlotE s S -> Index base 'SlotE -> ST s S
forall {kelem} a base (elem :: kelem) s.
Unbox a =>
MVector base elem s a -> Index base elem -> ST s a
C.readMV MVector base 'SlotE s S
mv (Int -> Index base 'SlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
j)
        if
          | proxy pol -> S -> Bool
forall (pol :: Pol) (proxy :: Pol -> *).
POL pol =>
proxy pol -> S -> Bool
forall (proxy :: Pol -> *). proxy pol -> S -> Bool
S.test proxy pol
pol S
w -> Int
-> Var base (PreImage pol 'EmptySlotE) -> ST s (MaybeFound base)
go (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Var base (PreImage pol 'EmptySlotE)
toSkip
          | Var base (PreImage pol 'EmptySlotE)
0 Var base (PreImage pol 'EmptySlotE)
-> Var base (PreImage pol 'EmptySlotE) -> Bool
forall a. Eq a => a -> a -> Bool
== Var base (PreImage pol 'EmptySlotE)
toSkip  -> MaybeFound base -> ST s (MaybeFound base)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MaybeFound base -> ST s (MaybeFound base))
-> MaybeFound base -> ST s (MaybeFound base)
forall a b. (a -> b) -> a -> b
$ Index base 'SlotE -> MaybeFound base
forall base. Index base 'SlotE -> MaybeFound base
JustFound (Int -> Index base 'SlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
j)
          | Bool
otherwise    -> Int
-> Var base (PreImage pol 'EmptySlotE) -> ST s (MaybeFound base)
go (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Var base (PreImage pol 'EmptySlotE)
toSkip Var base (PreImage pol 'EmptySlotE)
-> Var base (PreImage pol 'EmptySlotE)
-> Var base (PreImage pol 'EmptySlotE)
forall a. Num a => a -> a -> a
- Var base (PreImage pol 'EmptySlotE)
1)

findIthActiveInV ::
     C.Vector base SlotE S
  -> C.Index base ActiveSlotE
  -> MaybeFound base
findIthActiveInV :: forall base.
Vector base 'SlotE S -> Index base 'ActiveSlotE -> MaybeFound base
findIthActiveInV =
    Proxy 'Inverted
-> Vector base 'SlotE S
-> Index base (PreImage 'Inverted 'EmptySlotE)
-> MaybeFound base
forall (pol :: Pol) (proxy :: Pol -> *) base.
POL pol =>
proxy pol
-> Vector base 'SlotE S
-> Index base (PreImage pol 'EmptySlotE)
-> MaybeFound base
findIthEmptyInV Proxy 'Inverted
S.inverted

-----

-- | Trivial wrapper around 'countActivesInMV'
countActivesInV ::
     POL   pol
  => proxy pol
  -> C.Vector base SlotE S
  -> C.Size base (PreImage pol ActiveSlotE)
countActivesInV :: forall (pol :: Pol) (proxy :: Pol -> *) base.
POL pol =>
proxy pol
-> Vector base 'SlotE S -> Size base (PreImage pol 'ActiveSlotE)
countActivesInV proxy pol
pol Vector base 'SlotE S
v =
    Var base (PreImage pol 'ActiveSlotE)
-> Size base (PreImage pol 'ActiveSlotE)
forall {kelem} base (elem :: kelem).
Var base elem -> Size base elem
C.toSize (Var base (PreImage pol 'ActiveSlotE)
 -> Size base (PreImage pol 'ActiveSlotE))
-> Var base (PreImage pol 'ActiveSlotE)
-> Size base (PreImage pol 'ActiveSlotE)
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (Var base (PreImage pol 'ActiveSlotE)))
-> Var base (PreImage pol 'ActiveSlotE)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Var base (PreImage pol 'ActiveSlotE)))
 -> Var base (PreImage pol 'ActiveSlotE))
-> (forall s. ST s (Var base (PreImage pol 'ActiveSlotE)))
-> Var base (PreImage pol 'ActiveSlotE)
forall a b. (a -> b) -> a -> b
$ Vector base 'SlotE S -> ST s (MVector base 'SlotE s S)
forall {k1} {k2} a (base :: k1) (elem :: k2) s.
Unbox a =>
Vector base elem a -> ST s (MVector base elem s a)
C.unsafeThawV Vector base 'SlotE S
v ST s (MVector base 'SlotE s S)
-> (MVector base 'SlotE s S
    -> ST s (Var base (PreImage pol 'ActiveSlotE)))
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MVector base 'SlotE s S
mv -> proxy pol
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
countActivesInMV proxy pol
pol MVector base 'SlotE s S
mv

-- | The number of active slots in the vector
countActivesInMV ::
     POL   pol
  => proxy pol
  -> C.MVector base SlotE s S
  -> ST s (C.Var base (PreImage pol ActiveSlotE))
countActivesInMV :: forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
countActivesInMV proxy pol
pol MVector base 'SlotE s S
mv =
    (Var base (PreImage pol 'ActiveSlotE)
 -> S -> Var base (PreImage pol 'ActiveSlotE))
-> Var base (PreImage pol 'ActiveSlotE)
-> MVector (PrimState (ST s)) S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
(b -> a -> b) -> b -> MVector (PrimState m) a -> m b
MV.foldl'
        (\Var base (PreImage pol 'ActiveSlotE)
acc S
w -> if proxy pol -> S -> Bool
forall (pol :: Pol) (proxy :: Pol -> *).
POL pol =>
proxy pol -> S -> Bool
forall (proxy :: Pol -> *). proxy pol -> S -> Bool
S.test proxy pol
pol S
w then Var base (PreImage pol 'ActiveSlotE)
acc Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a. Num a => a -> a -> a
+ Var base (PreImage pol 'ActiveSlotE)
1 else Var base (PreImage pol 'ActiveSlotE)
acc)
        Var base (PreImage pol 'ActiveSlotE)
0
        MVector s S
MVector (PrimState (ST s)) S
mv'
  where
    C.MVector MVector s S
mv' = MVector base 'SlotE s S
mv

-----

-- | A density of active slots in a given window
--
-- @pol@ is the polarity to use for the active slots
--
-- TODO: rename to SomeDensity
data SomeDensityWindow pol =
  forall slidingWindow.
    SomeDensityWindow
        !(C.Var  slidingWindow (PreImage pol ActiveSlotE)) -- ^ Numerator: The active slots
        !(C.Size slidingWindow SlotE)                      -- ^ Denominator: The total amount of slots

instance Eq (SomeDensityWindow pol) where
    SomeDensityWindow Var slidingWindow (PreImage pol 'ActiveSlotE)
l1 Size slidingWindow 'SlotE
l2 == :: SomeDensityWindow pol -> SomeDensityWindow pol -> Bool
== SomeDensityWindow Var slidingWindow (PreImage pol 'ActiveSlotE)
r1 Size slidingWindow 'SlotE
r2 =
      Var slidingWindow (PreImage pol 'ActiveSlotE)
-> Forgotten (Count () (PreImage pol 'ActiveSlotE) Other)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Forgotten (Count () elem which)
C.forgetBase Var slidingWindow (PreImage pol 'ActiveSlotE)
l1 Forgotten (Count () (PreImage pol 'ActiveSlotE) Other)
-> Forgotten (Count () (PreImage pol 'ActiveSlotE) Other) -> Bool
forall a. Eq a => a -> a -> Bool
== Var slidingWindow (PreImage pol 'ActiveSlotE)
-> Forgotten (Count () (PreImage pol 'ActiveSlotE) Other)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Forgotten (Count () elem which)
C.forgetBase Var slidingWindow (PreImage pol 'ActiveSlotE)
r1 Bool -> Bool -> Bool
&& Size slidingWindow 'SlotE -> Forgotten (Count () 'SlotE Total)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Forgotten (Count () elem which)
C.forgetBase Size slidingWindow 'SlotE
l2 Forgotten (Count () 'SlotE Total)
-> Forgotten (Count () 'SlotE Total) -> Bool
forall a. Eq a => a -> a -> Bool
== Size slidingWindow 'SlotE -> Forgotten (Count () 'SlotE Total)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Forgotten (Count () elem which)
C.forgetBase Size slidingWindow 'SlotE
r2

instance Show (SomeDensityWindow pol) where
    showsPrec :: Int -> SomeDensityWindow pol -> ShowS
showsPrec Int
p (SomeDensityWindow Var slidingWindow (PreImage pol 'ActiveSlotE)
numer Size slidingWindow 'SlotE
denom) =
        Int -> ShowBuilder (SomeDensityWindow pol) -> ShowS
forall a.
NoFun "runShowsPrec" a (AbsError "runShowsPrec" a) =>
Int -> ShowBuilder a -> ShowS
Some.runShowsPrec Int
p
      (ShowBuilder (SomeDensityWindow pol) -> ShowS)
-> ShowBuilder (SomeDensityWindow pol) -> ShowS
forall a b. (a -> b) -> a -> b
$ (Var slidingWindow (PreImage pol 'ActiveSlotE)
 -> Size slidingWindow 'SlotE -> SomeDensityWindow pol)
-> String
-> ShowBuilder
     (Var slidingWindow (PreImage pol 'ActiveSlotE)
      -> Size slidingWindow 'SlotE -> SomeDensityWindow pol)
forall a. a -> String -> ShowBuilder a
Some.showCtor (forall (pol :: Pol) slidingWindow.
Var slidingWindow (PreImage pol 'ActiveSlotE)
-> Size slidingWindow 'SlotE -> SomeDensityWindow pol
SomeDensityWindow @pol) String
"SomeDensityWindow"
            ShowBuilder
  (Var slidingWindow (PreImage pol 'ActiveSlotE)
   -> Size slidingWindow 'SlotE -> SomeDensityWindow pol)
-> Var slidingWindow (PreImage pol 'ActiveSlotE)
-> ShowBuilder (Size slidingWindow 'SlotE -> SomeDensityWindow pol)
forall a b. Show a => ShowBuilder (a -> b) -> a -> ShowBuilder b
`Some.showArg` Var slidingWindow (PreImage pol 'ActiveSlotE)
numer
            ShowBuilder (Size slidingWindow 'SlotE -> SomeDensityWindow pol)
-> Size slidingWindow 'SlotE -> ShowBuilder (SomeDensityWindow pol)
forall a b. Show a => ShowBuilder (a -> b) -> a -> ShowBuilder b
`Some.showArg` Size slidingWindow 'SlotE
denom

instance Read (SomeDensityWindow pol) where
    readPrec :: ReadPrec (SomeDensityWindow pol)
readPrec =
        ReadBuilder (SomeDensityWindow pol)
-> ReadPrec (SomeDensityWindow pol)
forall a.
NoFun "runReadPrec" a (AbsError "runReadPrec" a) =>
ReadBuilder a -> ReadPrec a
Some.runReadPrec
      (ReadBuilder (SomeDensityWindow pol)
 -> ReadPrec (SomeDensityWindow pol))
-> ReadBuilder (SomeDensityWindow pol)
-> ReadPrec (SomeDensityWindow pol)
forall a b. (a -> b) -> a -> b
$ (Var Any (PreImage pol 'ActiveSlotE)
 -> Size Any 'SlotE -> SomeDensityWindow pol)
-> String
-> ReadBuilder
     (Var Any (PreImage pol 'ActiveSlotE)
      -> Size Any 'SlotE -> SomeDensityWindow pol)
forall a. a -> String -> ReadBuilder a
Some.readCtor Var Any (PreImage pol 'ActiveSlotE)
-> Size Any 'SlotE -> SomeDensityWindow pol
forall (pol :: Pol) slidingWindow.
Var slidingWindow (PreImage pol 'ActiveSlotE)
-> Size slidingWindow 'SlotE -> SomeDensityWindow pol
SomeDensityWindow String
"SomeDensityWindow"
            ReadBuilder
  (Var Any (PreImage pol 'ActiveSlotE)
   -> Size Any 'SlotE -> SomeDensityWindow pol)
-> ReadBuilder (Var Any (PreImage pol 'ActiveSlotE))
-> ReadBuilder (Size Any 'SlotE -> SomeDensityWindow pol)
forall a b. ReadBuilder (a -> b) -> ReadBuilder a -> ReadBuilder b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadBuilder (Var Any (PreImage pol 'ActiveSlotE))
forall a. Read a => ReadBuilder a
Some.readArg
            ReadBuilder (Size Any 'SlotE -> SomeDensityWindow pol)
-> ReadBuilder (Size Any 'SlotE)
-> ReadBuilder (SomeDensityWindow pol)
forall a b. ReadBuilder (a -> b) -> ReadBuilder a -> ReadBuilder b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadBuilder (Size Any 'SlotE)
forall a. Read a => ReadBuilder a
Some.readArg

-- | @fillInWindow pol (SomeDensityWindow k s) g mv@ mutates @mv@ to ensure
-- that the vector @take s $ mv ++ repeat (mkActive pol)@ has at least @k@
-- slots polarizely active.
--
-- Preconditions:
--
-- > lengthMV mv <= s
-- > k <= s
--
fillInWindow ::
  forall proxy pol base g s.
     (POL pol, R.StatefulGen g (ST s))
  => proxy pol
  -> SomeDensityWindow pol
  -> g
  -> C.MVector base SlotE s S
  -> ST s (C.Var base (PreImage pol ActiveSlotE))   -- ^ the count after filling
fillInWindow :: forall (proxy :: Pol -> *) (pol :: Pol) base g s.
(POL pol, StatefulGen g (ST s)) =>
proxy pol
-> SomeDensityWindow pol
-> g
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
fillInWindow proxy pol
pol (SomeDensityWindow Var slidingWindow (PreImage pol 'ActiveSlotE)
k Size slidingWindow 'SlotE
s) g
g MVector base 'SlotE s S
mv
        | Bool -> Bool
not (Var slidingWindow (PreImage pol 'ActiveSlotE) -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Var slidingWindow (PreImage pol 'ActiveSlotE)
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Size slidingWindow 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size slidingWindow 'SlotE
s) =
            String -> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a. HasCallStack => String -> a
error (String -> ST s (Var base (PreImage pol 'ActiveSlotE)))
-> String -> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a b. (a -> b) -> a -> b
$ String
"fillInWindow: assertion failure: k <= s: "
                    String -> ShowS
forall a. [a] -> [a] -> [a]
++ Var slidingWindow (PreImage pol 'ActiveSlotE) -> String
forall a. Show a => a -> String
show Var slidingWindow (PreImage pol 'ActiveSlotE)
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" <= " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Size slidingWindow 'SlotE -> String
forall a. Show a => a -> String
show Size slidingWindow 'SlotE
s
        | Bool -> Bool
not (Size base 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size base 'SlotE
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Size slidingWindow 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size slidingWindow 'SlotE
s) =
            String -> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a. HasCallStack => String -> a
error (String -> ST s (Var base (PreImage pol 'ActiveSlotE)))
-> String -> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a b. (a -> b) -> a -> b
$ String
"fillInWindow: assertion failure: sz <= s: "
                    String -> ShowS
forall a. [a] -> [a] -> [a]
++ Size base 'SlotE -> String
forall a. Show a => a -> String
show Size base 'SlotE
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" <= " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Size slidingWindow 'SlotE -> String
forall a. Show a => a -> String
show Size slidingWindow 'SlotE
s
        | Bool
otherwise = do
    -- how many active polarized slots @actual@ currently has
    Var base (PreImage pol 'ActiveSlotE)
initialActives <- proxy pol
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> ST s (Var base (PreImage pol 'ActiveSlotE))
countActivesInMV proxy pol
pol MVector base 'SlotE s S
mv


    -- discount the numerator accordingly if @mv@ is smaller than @s@
    --
    -- EG when a full-size @mv@ would reach past the 'Len'.
    --
    -- This discount reflects that we (very conservatively!) assume every
    -- truncated slot would be an active polarized slot.
    let discountedK :: C.Var base (PreImage pol ActiveSlotE)
        discountedK :: Var base (PreImage pol 'ActiveSlotE)
discountedK = Int -> Var base (PreImage pol 'ActiveSlotE)
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count (Int -> Var base (PreImage pol 'ActiveSlotE))
-> Int -> Var base (PreImage pol 'ActiveSlotE)
forall a b. (a -> b) -> a -> b
$ Var slidingWindow (PreImage pol 'ActiveSlotE) -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Var slidingWindow (PreImage pol 'ActiveSlotE)
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Size slidingWindow 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size slidingWindow 'SlotE
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Size base 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size base 'SlotE
sz)

    -- how many active polarized slots need to be added to @mv@
    let adding :: Var base (PreImage pol 'ActiveSlotE)
adding = Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a. Ord a => a -> a -> a
max Var base (PreImage pol 'ActiveSlotE)
0 (Var base (PreImage pol 'ActiveSlotE)
 -> Var base (PreImage pol 'ActiveSlotE))
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a b. (a -> b) -> a -> b
$ Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Var base elem
C.toVar Var base (PreImage pol 'ActiveSlotE)
discountedK Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a. Num a => a -> a -> a
- Var base (PreImage pol 'ActiveSlotE)
initialActives :: C.Var base (PreImage pol ActiveSlotE)

    -- draw from the empty polarized slots uniformly without replacement, a la Fisher-Yates shuffle
    Size base (PreImage pol 'ActiveSlotE)
-> (Index base (PreImage pol 'ActiveSlotE) -> ST s ()) -> ST s ()
forall {kelem} (f :: * -> *) base (elem :: kelem) a.
Applicative f =>
Size base elem -> (Index base elem -> f a) -> f ()
C.forRange_ (Var base (PreImage pol 'ActiveSlotE)
-> Size base (PreImage pol 'ActiveSlotE)
forall {kelem} base (elem :: kelem).
Var base elem -> Size base elem
C.toSize Var base (PreImage pol 'ActiveSlotE)
adding) ((Index base (PreImage pol 'ActiveSlotE) -> ST s ()) -> ST s ())
-> (Index base (PreImage pol 'ActiveSlotE) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Index base (PreImage pol 'ActiveSlotE)
alreadyAdded -> do
        let currentActives :: Size base (PreImage pol 'ActiveSlotE)
currentActives = Var base (PreImage pol 'ActiveSlotE)
-> Size base (PreImage pol 'ActiveSlotE)
forall {kelem} base (elem :: kelem).
Var base elem -> Size base elem
C.toSize (Var base (PreImage pol 'ActiveSlotE)
 -> Size base (PreImage pol 'ActiveSlotE))
-> Var base (PreImage pol 'ActiveSlotE)
-> Size base (PreImage pol 'ActiveSlotE)
forall a b. (a -> b) -> a -> b
$ Var base (PreImage pol 'ActiveSlotE)
initialActives Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a. Num a => a -> a -> a
+ Index base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Var base elem
C.toVar Index base (PreImage pol 'ActiveSlotE)
alreadyAdded
            currentEmpties :: Count base (PreImage pol 'EmptySlotE) Total
currentEmpties = proxy pol
-> Size base 'SlotE
-> Size base (PreImage pol 'ActiveSlotE)
-> Count base (PreImage pol 'EmptySlotE) Total
forall (proxy :: Pol -> *) (pol :: Pol) base which.
proxy pol
-> Size base 'SlotE
-> Count base (PreImage pol 'ActiveSlotE) which
-> Count base (PreImage pol 'EmptySlotE) which
S.complementActive proxy pol
pol Size base 'SlotE
sz Size base (PreImage pol 'ActiveSlotE)
currentActives

        Index base (PreImage pol 'EmptySlotE)
whichEmptyToFlip <- Count base (PreImage pol 'EmptySlotE) Total
-> g -> ST s (Index base (PreImage pol 'EmptySlotE))
forall {kelem} g (m :: * -> *) base (elem :: kelem).
StatefulGen g m =>
Size base elem -> g -> m (Index base elem)
C.uniformIndex Count base (PreImage pol 'EmptySlotE) Total
currentEmpties g
g

        Index base 'SlotE
slot <- proxy pol
-> MVector base 'SlotE s S
-> Index base (PreImage pol 'EmptySlotE)
-> ST s (MaybeFound base)
forall (proxy :: Pol -> *) (pol :: Pol) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S
-> Index base (PreImage pol 'EmptySlotE)
-> ST s (MaybeFound base)
findIthEmptyInMV proxy pol
pol MVector base 'SlotE s S
mv Index base (PreImage pol 'EmptySlotE)
whichEmptyToFlip ST s (MaybeFound base)
-> (MaybeFound base -> Index base 'SlotE)
-> ST s (Index base 'SlotE)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
            JustFound Index base 'SlotE
i  -> Index base 'SlotE
i
            MaybeFound base
NothingFound -> String -> Index base 'SlotE
forall a. HasCallStack => String -> a
error String
"impossible! fillInWindow"

        proxy pol
-> MVector base 'SlotE s S -> Index base 'SlotE -> ST s ()
forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S -> Index base 'SlotE -> ST s ()
setMV proxy pol
pol MVector base 'SlotE s S
mv Index base 'SlotE
slot

    Var base (PreImage pol 'ActiveSlotE)
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Var base (PreImage pol 'ActiveSlotE)
 -> ST s (Var base (PreImage pol 'ActiveSlotE)))
-> Var base (PreImage pol 'ActiveSlotE)
-> ST s (Var base (PreImage pol 'ActiveSlotE))
forall a b. (a -> b) -> a -> b
$ Var base (PreImage pol 'ActiveSlotE)
initialActives Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
-> Var base (PreImage pol 'ActiveSlotE)
forall a. Num a => a -> a -> a
+ Var base (PreImage pol 'ActiveSlotE)
adding
  where
    sz :: Size base 'SlotE
sz = MVector base 'SlotE s S -> Size base 'SlotE
forall {kelem} a base (elem :: kelem) s.
Unbox a =>
MVector base elem s a -> Size base elem
C.lengthMV MVector base 'SlotE s S
mv :: C.Size base SlotE

-----

testV :: POL pol => proxy pol -> C.Vector base SlotE S -> C.Index base SlotE -> Bool
testV :: forall (pol :: Pol) (proxy :: Pol -> *) base.
POL pol =>
proxy pol -> Vector base 'SlotE S -> Index base 'SlotE -> Bool
testV proxy pol
pol Vector base 'SlotE S
mv Index base 'SlotE
i = proxy pol -> S -> Bool
forall (pol :: Pol) (proxy :: Pol -> *).
POL pol =>
proxy pol -> S -> Bool
forall (proxy :: Pol -> *). proxy pol -> S -> Bool
S.test proxy pol
pol (Vector base 'SlotE S -> Index base 'SlotE -> S
forall {kelem} a base (elem :: kelem).
Unbox a =>
Vector base elem a -> Index base elem -> a
C.readV Vector base 'SlotE S
mv Index base 'SlotE
i)

testMV :: POL pol => proxy pol -> C.MVector base SlotE s S -> C.Index base SlotE -> ST s Bool
testMV :: forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S -> Index base 'SlotE -> ST s Bool
testMV proxy pol
pol MVector base 'SlotE s S
mv Index base 'SlotE
i = do
     S
w <- MVector base 'SlotE s S -> Index base 'SlotE -> ST s S
forall {kelem} a base (elem :: kelem) s.
Unbox a =>
MVector base elem s a -> Index base elem -> ST s a
C.readMV MVector base 'SlotE s S
mv Index base 'SlotE
i
     Bool -> ST s Bool
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$ proxy pol -> S -> Bool
forall (pol :: Pol) (proxy :: Pol -> *).
POL pol =>
proxy pol -> S -> Bool
forall (proxy :: Pol -> *). proxy pol -> S -> Bool
S.test proxy pol
pol S
w

setMV :: POL pol => proxy pol -> C.MVector base SlotE s S -> C.Index base SlotE -> ST s ()
setMV :: forall (pol :: Pol) (proxy :: Pol -> *) base s.
POL pol =>
proxy pol
-> MVector base 'SlotE s S -> Index base 'SlotE -> ST s ()
setMV proxy pol
pol MVector base 'SlotE s S
mv Index base 'SlotE
i = MVector base 'SlotE s S -> Index base 'SlotE -> S -> ST s ()
forall {kelem} a base (elem :: kelem) s.
Unbox a =>
MVector base elem s a -> Index base elem -> a -> ST s ()
C.writeMV MVector base 'SlotE s S
mv Index base 'SlotE
i (S -> ST s ()) -> S -> ST s ()
forall a b. (a -> b) -> a -> b
$ proxy pol -> S
forall (pol :: Pol) (proxy :: Pol -> *). POL pol => proxy pol -> S
forall (proxy :: Pol -> *). proxy pol -> S
S.mkActive proxy pol
pol