File: NFreedom.hs

package info (click to toggle)
haskell-arithmoi 0.13.2.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 964 kB
  • sloc: haskell: 10,379; makefile: 3
file content (161 lines) | stat: -rw-r--r-- 5,954 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
-- |
-- 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 _ _ 0 = U.empty
sieveBlockNFree n lowIndex len'
  = runST $ do
    as <- MU.replicate (wordToInt len') True
    forM_ ps $ \p -> do
      let pPow :: a
          pPow = p ^ n
          offset :: a
          offset = negate lowIndex `mod` 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 = [offset, offset + pPow .. len - 1]
      forM_ indices $ \ix ->
          MU.write as (fromIntegral ix) False
    U.freeze as

  where
    len :: a
    len = fromIntegral len'

    highIndex :: a
    highIndex = lowIndex + len - 1

    ps :: [a]
    ps = if highIndex < 4 then [] else map unPrime [nextPrime 2 .. precPrime (integerSquareRoot 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 0 = [1]
nFrees 1 = [1]
nFrees n = concatMap (uncurry (nFreesBlock n)) $ zip bounds 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 = take 55 (iterate (2 *) 256) ++ repeat (intToWord (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 = scanl' (+) 1 $ map fromIntegral 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 0 lo _ = help lo
nFreesBlock 1 lo _ = help lo
nFreesBlock n lowIndex 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' = wordToInt len
        -- * Twice, immediately below, to create the range of indices with
        -- which to query @bs@.
        len'' :: a
        len'' = fromIntegral len
        bs  = sieveBlockNFree n lowIndex len
    in map snd .
       filter ((bs U.!) . fst) .
       zip [0 .. len' - 1] $ [lowIndex .. lowIndex + len'']
{-# INLINE nFreesBlock #-}

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