-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.NFreedom
-- Copyright:   (c) 2018 Alexandre Rodrigues Baldé
-- Licence:     MIT
-- Maintainer:  Alexandre Rodrigues Baldé <alexandrer_b@outlook.com>
--
-- N-free number generation.
--

{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.ArithmeticFunctions.NFreedom
  ( nFrees
  , nFreesBlock
  , sieveBlockNFree
  ) where

import Control.Monad                         (forM_)
import Control.Monad.ST                      (runST)
import Data.Bits (Bits)
import Data.List                             (scanl')
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector.Unboxed.Mutable as MU

import Math.NumberTheory.Roots
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils.FromIntegral

-- | Evaluate the `Math.NumberTheory.ArithmeticFunctions.isNFree` function over a block.
-- Value at @0@, if zero falls into block, is undefined.
--
-- This function should __**not**__ be used with a negative lower bound.
-- If it is, the result is undefined.
-- Furthermore, do not:
--
-- * use a block length greater than @maxBound :: Int@.
-- * use a power that is either of @0, 1@.
--
-- None of these preconditions are checked, and if any occurs, the result is
-- undefined, __if__ the function terminates.
--
-- >>> sieveBlockNFree 2 1 10
-- [True,True,True,False,True,True,True,False,False,True]
sieveBlockNFree
  :: forall a. (Integral a, Enum (Prime a), Bits a, UniqueFactorisation a)
  => Word
  -- ^ Power whose @n@-freedom will be checked.
  -> a
  -- ^ Lower index of the block.
  -> Word
  -- ^ Length of the block.
  -> U.Vector Bool
  -- ^ Vector of flags, where @True@ at index @i@ means the @i@-th element of
  -- the block is @n@-free.
sieveBlockNFree :: Word -> a -> Word -> Vector Bool
sieveBlockNFree Word
_ a
_ Word
0 = Vector Bool
forall a. Unbox a => Vector a
U.empty
sieveBlockNFree Word
n a
lowIndex Word
len'
  = (forall s. ST s (Vector Bool)) -> Vector Bool
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bool)) -> Vector Bool)
-> (forall s. ST s (Vector Bool)) -> Vector Bool
forall a b. (a -> b) -> a -> b
$ do
    MVector s Bool
as <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate (Word -> Int
wordToInt Word
len') Bool
True
    [a] -> (a -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
ps ((a -> ST s ()) -> ST s ()) -> (a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \a
p -> do
      let pPow :: a
          pPow :: a
pPow = a
p a -> Word -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
n
          offset :: a
          offset :: a
offset = a -> a
forall a. Num a => a -> a
negate a
lowIndex a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
pPow
          -- The second argument in @Data.Vector.Unboxed.Mutable.write@ is an
          -- @Int@, so to avoid segmentation faults or out-of-bounds errors,
          -- the enumeration's higher bound must always be less than
          -- @maxBound :: Int@.
          -- As mentioned above, this is not checked when using this function
          -- by itself, but is carefully managed when this function is called
          -- by @nFrees@, see the comments in it.
          indices :: [a]
          indices :: [a]
indices = [a
offset, a
offset a -> a -> a
forall a. Num a => a -> a -> a
+ a
pPow .. a
len a -> a -> a
forall a. Num a => a -> a -> a
- a
1]
      [a] -> (a -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
indices ((a -> ST s ()) -> ST s ()) -> (a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \a
ix ->
          MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.write MVector s Bool
MVector (PrimState (ST s)) Bool
as (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
ix) Bool
False
    MVector (PrimState (ST s)) Bool -> ST s (Vector Bool)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.freeze MVector s Bool
MVector (PrimState (ST s)) Bool
as

  where
    len :: a
    len :: a
len = Word -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len'

    highIndex :: a
    highIndex :: a
highIndex = a
lowIndex a -> a -> a
forall a. Num a => a -> a -> a
+ a
len a -> a -> a
forall a. Num a => a -> a -> a
- a
1

    ps :: [a]
    ps :: [a]
ps = if a
highIndex a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
4 then [] else (Prime a -> a) -> [Prime a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Prime a -> a
forall a. Prime a -> a
unPrime [a -> Prime a
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime a
2 .. a -> Prime a
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
precPrime (a -> a
forall a. Integral a => a -> a
integerSquareRoot a
highIndex)]

-- | For a given nonnegative integer power @n@, generate all @n@-free
-- numbers in ascending order, starting at @1@.
--
-- When @n@ is @0@ or @1@, the resulting list is @[1]@.
nFrees
    :: forall a. (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
    => Word
    -- ^ Power @n@ to be used to generate @n@-free numbers.
    -> [a]
    -- ^ Generated infinite list of @n@-free numbers.
nFrees :: Word -> [a]
nFrees Word
0 = [a
1]
nFrees Word
1 = [a
1]
nFrees Word
n = ((a, Word) -> [a]) -> [(a, Word)] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((a -> Word -> [a]) -> (a, Word) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Word -> a -> Word -> [a]
forall a.
(Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
Word -> a -> Word -> [a]
nFreesBlock Word
n)) ([(a, Word)] -> [a]) -> [(a, Word)] -> [a]
forall a b. (a -> b) -> a -> b
$ [a] -> [Word] -> [(a, Word)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
bounds [Word]
strides
  where
    -- The 56th element of @iterate (2 *) 256@ is @2^64 :: Word == 0@, so to
    -- avoid overflow only the first 55 elements of this list are used.
    -- After those, since @maxBound :: Int@ is the largest a vector can be,
    -- this value is just repeated. This means after a few dozen iterations,
    -- the sieve will stop increasing in size.
    strides :: [Word]
    strides :: [Word]
strides = Int -> [Word] -> [Word]
forall a. Int -> [a] -> [a]
take Int
55 ((Word -> Word) -> Word -> [Word]
forall a. (a -> a) -> a -> [a]
iterate (Word
2 Word -> Word -> Word
forall a. Num a => a -> a -> a
*) Word
256) [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ Word -> [Word]
forall a. a -> [a]
repeat (Int -> Word
intToWord (Int
forall a. Bounded a => a
maxBound :: Int))

    -- Infinite list of lower bounds at which @sieveBlockNFree@ will be
    -- applied. This has type @Integral a => a@ instead of @Word@ because
    -- unlike the sizes of the sieve that eventually stop increasing (see
    -- above comment), the lower bound at which @sieveBlockNFree@ is called does not.
    bounds :: [a]
    bounds :: [a]
bounds = (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
1 ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (Word -> a) -> [Word] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Word -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word]
strides

-- | Generate @n@-free numbers in a block starting at a certain value.
-- The length of the list is determined by the value passed in as the third
-- argument. It will be lesser than or equal to this value.
--
-- This function should not be used with a negative lower bound. If it is,
-- the result is undefined.
--
-- The block length cannot exceed @maxBound :: Int@, this precondition is not
-- checked.
--
-- As with @nFrees@, passing @n = 0, 1@ results in an empty list.
nFreesBlock
    :: forall a . (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
    => Word
    -- ^ Power @n@ to be used to generate @n@-free numbers.
    -> a
    -- ^ Starting number in the block.
    -> Word
    -- ^ Maximum length of the block to be generated.
    -> [a]
    -- ^ Generated list of @n@-free numbers.
nFreesBlock :: Word -> a -> Word -> [a]
nFreesBlock Word
0 a
lo Word
_ = a -> [a]
forall a. Integral a => a -> [a]
help a
lo
nFreesBlock Word
1 a
lo Word
_ = a -> [a]
forall a. Integral a => a -> [a]
help a
lo
nFreesBlock Word
n a
lowIndex Word
len =
    let -- When indexing the array of flags @bs@, the index has to be an
        -- @Int@. As such, it's necessary to cast @strd@ twice.
        -- * Once, immediately below, to create the range of values whose
        -- @n@-freedom will be tested. Since @nFrees@ has return type
        -- @[a]@, this cannot be avoided as @strides@ has type @[Word]@.
        len' :: Int
        len' :: Int
len' = Word -> Int
wordToInt Word
len
        -- * Twice, immediately below, to create the range of indices with
        -- which to query @bs@.
        len'' :: a
        len'' :: a
len'' = Word -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len
        bs :: Vector Bool
bs  = Word -> a -> Word -> Vector Bool
forall a.
(Integral a, Enum (Prime a), Bits a, UniqueFactorisation a) =>
Word -> a -> Word -> Vector Bool
sieveBlockNFree Word
n a
lowIndex Word
len
    in ((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd ([(Int, a)] -> [a]) -> ([a] -> [(Int, a)]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
       ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Vector Bool
bs Vector Bool -> Int -> Bool
forall a. Unbox a => Vector a -> Int -> a
U.!) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) ([(Int, a)] -> [(Int, a)])
-> ([a] -> [(Int, a)]) -> [a] -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
       [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 .. Int
len' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ [a
lowIndex .. a
lowIndex a -> a -> a
forall a. Num a => a -> a -> a
+ a
len'']
{-# INLINE nFreesBlock #-}

help :: Integral a => a -> [a]
help :: a -> [a]
help a
1 = [a
1]
help a
_ = []
{-# INLINE help #-}