-- |
-- Module      : Data.Vector.Fusion.Bundle.Size
-- Copyright   : (c) Roman Leshchinskiy 2008-2010
--                   Alexey Kuleshevich 2020-2022
--                   Aleksey Khudyakov 2020-2022
--                   Andrew Lelechenko 2020-2022
-- License     : BSD-style
--
-- Maintainer  : Haskell Libraries Team <libraries@haskell.org>
-- Stability   : experimental
-- Portability : portable
--
-- Size hints for streams.
--

module Data.Vector.Fusion.Bundle.Size (
  Size(..), clampedSubtract, smaller, smallerThan, larger, toMax, upperBound, lowerBound
) where

import Data.Vector.Fusion.Util ( delay_inline )

-- | Size hint
data Size = Exact {-# UNPACK #-} !Int -- ^ Exact size
          | Max   {-# UNPACK #-} !Int -- ^ Upper bound on the size
          | Unknown                   -- ^ Unknown size
        deriving( Size -> Size -> Bool
(Size -> Size -> Bool) -> (Size -> Size -> Bool) -> Eq Size
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Size -> Size -> Bool
== :: Size -> Size -> Bool
$c/= :: Size -> Size -> Bool
/= :: Size -> Size -> Bool
Eq, Int -> Size -> ShowS
[Size] -> ShowS
Size -> String
(Int -> Size -> ShowS)
-> (Size -> String) -> ([Size] -> ShowS) -> Show Size
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Size -> ShowS
showsPrec :: Int -> Size -> ShowS
$cshow :: Size -> String
show :: Size -> String
$cshowList :: [Size] -> ShowS
showList :: [Size] -> ShowS
Show )

instance Num Size where
  Exact Int
m + :: Size -> Size -> Size
+ Exact Int
n = (Int -> Size) -> Int -> Int -> Size
checkedAdd Int -> Size
Exact Int
m Int
n
  Exact Int
m + Max   Int
n = (Int -> Size) -> Int -> Int -> Size
checkedAdd Int -> Size
Max Int
m Int
n

  Max   Int
m + Exact Int
n = (Int -> Size) -> Int -> Int -> Size
checkedAdd Int -> Size
Max Int
m Int
n
  Max   Int
m + Max   Int
n = (Int -> Size) -> Int -> Int -> Size
checkedAdd Int -> Size
Max Int
m Int
n

  Size
_       + Size
_       = Size
Unknown


  Exact Int
m - :: Size -> Size -> Size
- Exact Int
n = (Int -> Size) -> Int -> Int -> Size
checkedSubtract Int -> Size
Exact Int
m Int
n
  Exact Int
m - Max   Int
_ = Int -> Size
Max   Int
m

  Max   Int
m - Exact Int
n = (Int -> Size) -> Int -> Int -> Size
checkedSubtract Int -> Size
Max Int
m Int
n
  Max   Int
m - Max   Int
_ = Int -> Size
Max   Int
m
  Max   Int
m - Size
Unknown = Int -> Size
Max   Int
m

  Size
_       - Size
_       = Size
Unknown


  fromInteger :: Integer -> Size
fromInteger Integer
n     = Int -> Size
Exact (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n)

  * :: Size -> Size -> Size
(*)    = String -> Size -> Size -> Size
forall a. HasCallStack => String -> a
error String
"vector: internal error * for Bundle.size isn't defined"
  abs :: Size -> Size
abs    = String -> Size -> Size
forall a. HasCallStack => String -> a
error String
"vector: internal error abs for Bundle.size isn't defined"
  signum :: Size -> Size
signum = String -> Size -> Size
forall a. HasCallStack => String -> a
error String
"vector: internal error signum for Bundle.size isn't defined"

{-# INLINE checkedAdd #-}
checkedAdd :: (Int -> Size) -> Int -> Int -> Size
checkedAdd :: (Int -> Size) -> Int -> Int -> Size
checkedAdd Int -> Size
con Int
m Int
n
    -- Note: we assume m and n are >= 0.
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m Bool -> Bool -> Bool
|| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
      String -> Size
forall a. HasCallStack => String -> a
error (String -> Size) -> String -> Size
forall a b. (a -> b) -> a -> b
$ String
"Data.Vector.Fusion.Bundle.Size.checkedAdd: overflow: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
r
  | Bool
otherwise = Int -> Size
con Int
r
  where
    r :: Int
r = Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n

{-# INLINE checkedSubtract #-}
checkedSubtract :: (Int -> Size) -> Int -> Int -> Size
checkedSubtract :: (Int -> Size) -> Int -> Int -> Size
checkedSubtract Int -> Size
con Int
m Int
n
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
      String -> Size
forall a. HasCallStack => String -> a
error (String -> Size) -> String -> Size
forall a b. (a -> b) -> a -> b
$ String
"Data.Vector.Fusion.Bundle.Size.checkedSubtract: underflow: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
r
  | Bool
otherwise = Int -> Size
con Int
r
  where
    r :: Int
r = Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n

-- | Subtract two sizes with clamping to 0, for drop-like things
{-# INLINE clampedSubtract #-}
clampedSubtract :: Size -> Size -> Size
clampedSubtract :: Size -> Size -> Size
clampedSubtract (Exact Int
m) (Exact Int
n) = Int -> Size
Exact (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n))
clampedSubtract (Max   Int
m) (Exact Int
n)
  | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = Int -> Size
Exact Int
0
  | Bool
otherwise = Int -> Size
Max (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n)
clampedSubtract (Exact Int
m) (Max   Int
_) = Int -> Size
Max Int
m
clampedSubtract (Max   Int
m) (Max   Int
_) = Int -> Size
Max Int
m
clampedSubtract Size
_         Size
_ = Size
Unknown

-- | Minimum of two size hints
smaller :: Size -> Size -> Size
{-# INLINE smaller #-}
smaller :: Size -> Size -> Size
smaller (Exact Int
m) (Exact Int
n) = Int -> Size
Exact ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smaller (Exact Int
m) (Max   Int
n) = Int -> Size
Max   ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smaller (Exact Int
m) Size
Unknown   = Int -> Size
Max   Int
m
smaller (Max   Int
m) (Exact Int
n) = Int -> Size
Max   ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smaller (Max   Int
m) (Max   Int
n) = Int -> Size
Max   ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smaller (Max   Int
m) Size
Unknown   = Int -> Size
Max   Int
m
smaller Size
Unknown   (Exact Int
n) = Int -> Size
Max   Int
n
smaller Size
Unknown   (Max   Int
n) = Int -> Size
Max   Int
n
smaller Size
Unknown   Size
Unknown   = Size
Unknown

-- | Select a safe smaller than known size.
smallerThan :: Int -> Size -> Size
{-# INLINE smallerThan #-}
smallerThan :: Int -> Size -> Size
smallerThan Int
m (Exact Int
n) = Int -> Size
Exact ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smallerThan Int
m (Max   Int
n) = Int -> Size
Max   ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n)
smallerThan Int
_ Size
Unknown   = Size
Unknown


-- | Maximum of two size hints
larger :: Size -> Size -> Size
{-# INLINE larger #-}
larger :: Size -> Size -> Size
larger (Exact Int
m) (Exact Int
n)             = Int -> Size
Exact ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
m Int
n)
larger (Exact Int
m) (Max   Int
n) | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = Int -> Size
Exact Int
m
                           | Bool
otherwise = Int -> Size
Max   Int
n
larger (Max   Int
m) (Exact Int
n) | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
m    = Int -> Size
Exact Int
n
                           | Bool
otherwise = Int -> Size
Max   Int
m
larger (Max   Int
m) (Max   Int
n)             = Int -> Size
Max   ((Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
delay_inline Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
m Int
n)
larger Size
_         Size
_                     = Size
Unknown

-- | Convert a size hint to an upper bound
toMax :: Size -> Size
toMax :: Size -> Size
toMax (Exact Int
n) = Int -> Size
Max Int
n
toMax (Max   Int
n) = Int -> Size
Max Int
n
toMax Size
Unknown   = Size
Unknown

-- | Compute the minimum size from a size hint
lowerBound :: Size -> Int
lowerBound :: Size -> Int
lowerBound (Exact Int
n) = Int
n
lowerBound Size
_         = Int
0

-- | Compute the maximum size from a size hint if possible
upperBound :: Size -> Maybe Int
upperBound :: Size -> Maybe Int
upperBound (Exact Int
n) = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
upperBound (Max   Int
n) = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
upperBound Size
Unknown   = Maybe Int
forall a. Maybe a
Nothing