-- |
-- Module:      Math.NumberTheory.Utils.Hyperbola
-- Copyright:   (c) 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Highest points under hyperbola.
--

module Math.NumberTheory.Utils.Hyperbola
  ( pointsUnderHyperbola
  ) where

import Data.Bits

import Math.NumberTheory.Roots

-- | Straightforward computation of
-- [ n `quot` x | x <- [hi, hi - 1 .. lo] ].
-- Unfortunately, such list generator performs poor,
-- so we fall back to manual recursion.
pointsUnderHyperbola0 :: Int -> Int -> Int -> [Int]
pointsUnderHyperbola0 :: Int -> Int -> Int -> [Int]
pointsUnderHyperbola0 Int
n Int
lo Int
hi
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"pointsUnderHyperbola0: first argument must be non-negative"
  | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0   = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"pointsUnderHyperbola0: second argument must be positive"
  | Bool
otherwise = Int -> [Int]
go Int
hi
    where
      go :: Int -> [Int]
go Int
x
        | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
lo    = []
        | Bool
otherwise = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> [Int]
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

data Bresenham = Bresenham
  {  Bresenham -> Int
bresX       :: !Int
  ,  Bresenham -> Int
bresBeta    :: !Int
  , Bresenham -> Int
_bresGamma   :: !Int
  , Bresenham -> Int
_bresDelta1  :: !Int
  , Bresenham -> Int
_bresEpsilon :: !Int
  }

initBresenham :: Int -> Int -> Bresenham
initBresenham :: Int -> Int -> Bresenham
initBresenham Int
n Int
x = Int -> Int -> Int -> Int -> Int -> Bresenham
Bresenham Int
x Int
beta Int
gamma Int
delta1 Int
epsilon
  where
    beta :: Int
beta    = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
x
    epsilon :: Int
epsilon = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
x
    delta1 :: Int
delta1  = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
beta
    gamma :: Int
gamma   = Int
beta Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
delta1

-- | bresenham(x+1) -> bresenham(x) for x >= (2n)^1/3
stepBack :: Bresenham -> Bresenham
stepBack :: Bresenham -> Bresenham
stepBack (Bresenham Int
x' Int
beta' Int
gamma' Int
delta1' Int
epsilon')
  | Int
eps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 {- delta2 = 2 -}
  = let delta1 :: Int
delta1 = Int
delta1' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 in Int -> Int -> Int -> Int -> Int -> Bresenham
Bresenham Int
x (Int
beta' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1) (Int
gamma' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Int
delta1 (Int
eps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1)
  | Int
eps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x {- delta1 = 1 -}
  = let delta1 :: Int
delta1 = Int
delta1' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 in Int -> Int -> Int -> Int -> Int -> Bresenham
Bresenham Int
x (Int
beta' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1) (Int
gamma' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) Int
delta1 (Int
eps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x)
  | Int
eps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 {- delta2 =  0 -}
  = Int -> Int -> Int -> Int -> Int -> Bresenham
Bresenham Int
x (Int
beta' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1') (Int
gamma' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1' Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Int
delta1' Int
eps
  | Bool
otherwise {- delta2 = -1 -}
  = let delta1 :: Int
delta1 = Int
delta1' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 in Int -> Int -> Int -> Int -> Int -> Bresenham
Bresenham Int
x (Int
beta' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1) (Int
gamma' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) Int
delta1 (Int
eps Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
  where
    x :: Int
x       = Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    eps :: Int
eps     = Int
epsilon' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
gamma'
{-# INLINE stepBack #-}

-- | Division-free computation of
-- [ n `quot` x | x <- [hi, hi - 1 .. lo] ].
-- In other words, we compute y-coordinates of highest integral points
-- under hyperbola @x * y = n@ between @x = lo@ and @x = hi@ in reverse order.
--
-- The implementation follows section 5 of <https://arxiv.org/pdf/1206.3369.pdf A successive approximation algorithm for computing the divisor summatory function>
-- by R. Sladkey.
-- It is 2x faster than a trivial implementation for 'Int'.
pointsUnderHyperbola :: Int -> Int -> Int -> [Int]
pointsUnderHyperbola :: Int -> Int -> Int -> [Int]
pointsUnderHyperbola Int
n Int
lo Int
hi
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0        = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"pointsUnderHyperbola: first argument must be non-negative"
  | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0      = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"pointsUnderHyperbola: second argument must be positive"
  | Int
hi Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
lo     = []
  | Int
hi Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lo     = [Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
lo]
  | Bool
otherwise    = Bresenham -> [Int]
go (Int -> Int -> Bresenham
initBresenham Int
n Int
hi)
  where
    mid :: Int
mid = (Int -> Int
forall a. Integral a => a -> a
integerCubeRoot (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
lo
    go :: Bresenham -> [Int]
go Bresenham
h
      | Bresenham -> Int
bresX Bresenham
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mid = Int -> Int -> Int -> [Int]
pointsUnderHyperbola0 Int
n Int
lo ((Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`min` Int
hi)
      | Bool
otherwise = Bresenham -> Int
bresBeta Bresenham
h Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Bresenham -> [Int]
go (Bresenham -> Bresenham
stepBack Bresenham
h)