{-# LANGUAGE CPP #-}
#ifdef __GLASGOW_HASKELL__
{-# LANGUAGE DeriveLift #-}
{-# LANGUAGE StandaloneDeriving #-}
#endif

-- |
-- = WARNING
--
-- This module is considered __internal__.
--
-- The Package Versioning Policy __does not apply__.
--
-- The contents of this module may change __in any way whatsoever__
-- and __without any warning__ between minor versions of this package.
--
-- Authors importing this module are expected to track development
-- closely.
--
-- = Description
--
-- This module defines common constructs used by both "Data.IntSet" and
-- "Data.IntMap".
--
-- @since 0.8
--

module Data.IntSet.Internal.IntTreeCommons
  ( Key
  , Prefix(..)
  , nomatch
  , left
  , signBranch
  , TreeTreeBranch(..)
  , treeTreeBranch
  , mask
  , branchMask
  , i2w
  , Order(..)
  ) where

import Data.Bits (Bits(..), countLeadingZeros)
import Utils.Containers.Internal.BitUtil (wordSize)

#ifdef __GLASGOW_HASKELL__
import Language.Haskell.TH.Syntax (Lift)
-- See Note [ Template Haskell Dependencies ]
import Language.Haskell.TH ()
#endif


type Key = Int

-- | A @Prefix@ represents some prefix of high-order bits of an @Int@.
--
-- A @Prefix@ is usually considered in the context of a
-- 'Data.IntSet.Internal.Bin' or 'Data.IntMap.Internal.Bin'.

-- See Note [IntSet structure and invariants] in Data.IntSet.Internal and
-- Note [IntMap structure and invariants] in Data.IntMap.Internal for details.
newtype Prefix = Prefix { Prefix -> Int
unPrefix :: Int }
  deriving Prefix -> Prefix -> Bool
(Prefix -> Prefix -> Bool)
-> (Prefix -> Prefix -> Bool) -> Eq Prefix
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Prefix -> Prefix -> Bool
== :: Prefix -> Prefix -> Bool
$c/= :: Prefix -> Prefix -> Bool
/= :: Prefix -> Prefix -> Bool
Eq

#ifdef __GLASGOW_HASKELL__
deriving instance Lift Prefix
#endif

-- | Whether the @Int@ does not start with the given @Prefix@.
--
-- An @Int@ starts with a @Prefix@ if it shares the high bits with the internal
-- @Int@ value of the @Prefix@ up to the mask bit.
--
-- @nomatch@ is usually used to determine whether a key belongs in a @Bin@,
-- since all keys in a @Bin@ share a @Prefix@.
nomatch :: Int -> Prefix -> Bool
nomatch :: Int -> Prefix -> Bool
nomatch Int
i Prefix
p = (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
px) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
prefixMask Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
  where
    px :: Int
px = Prefix -> Int
unPrefix Prefix
p
    prefixMask :: Int
prefixMask = Int
px Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` (-Int
px)
{-# INLINE nomatch #-}

-- | Whether the @Int@ is to the left of the split created by a @Bin@ with this
-- @Prefix@.
--
-- This does not imply that the @Int@ belongs in this @Bin@. That fact is
-- usually determined first using @nomatch@.
left :: Int -> Prefix -> Bool
left :: Int -> Prefix -> Bool
left Int
i Prefix
p = Int -> Word
i2w Int
i Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Word
i2w (Prefix -> Int
unPrefix Prefix
p)
{-# INLINE left #-}

-- | A @TreeTreeBranch@ is returned by 'treeTreeBranch' to indicate how two
-- @Bin@s relate to each other.
--
-- Consider that @A@ and @B@ are the @Bin@s whose @Prefix@es are given to
-- @treeTreeBranch@ as the first and second arguments respectively.
data TreeTreeBranch
  = ABL  -- ^ A contains B in the left child
  | ABR  -- ^ A contains B in the right child
  | BAL  -- ^ B contains A in the left child
  | BAR  -- ^ B contains A in the right child
  | EQL  -- ^ A and B have equal prefixes
  | NOM  -- ^ A and B have prefixes that do not match

-- | Calculates how two @Bin@s relate to each other by comparing their
-- @Prefix@es.

-- Notes:
-- * pw .|. (pw-1) sets every bit below the mask bit to 1. This is the greatest
--   key the Bin can have.
-- * pw .&. (pw-1) sets the mask bit and every bit below it to 0. This is the
--   smallest key the Bin can have.
--
-- First, we compare the prefixes to each other. Then we compare a prefix
-- against the greatest/smallest keys the other prefix's Bin could have. This is
-- enough to determine how the two Bins relate to each other. The conditions can
-- be stated as:
--
-- * If pw1 from Bin A is less than pw2 from Bin B, and pw2 is <= the greatest
--   key of Bin A, then Bin A contains Bin B in its right child.
-- * ...and so on

treeTreeBranch :: Prefix -> Prefix -> TreeTreeBranch
treeTreeBranch :: Prefix -> Prefix -> TreeTreeBranch
treeTreeBranch Prefix
p1 Prefix
p2 = case Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Word
pw1 Word
pw2 of
  Ordering
LT | Word
pw2 Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word -> Word
forall {a}. (Bits a, Num a) => a -> a
greatest Word
pw1 -> TreeTreeBranch
ABR
     | Word -> Word
forall {a}. (Bits a, Num a) => a -> a
smallest Word
pw2 Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
pw1 -> TreeTreeBranch
BAL
     | Bool
otherwise           -> TreeTreeBranch
NOM
  Ordering
GT | Word
pw1 Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word -> Word
forall {a}. (Bits a, Num a) => a -> a
greatest Word
pw2 -> TreeTreeBranch
BAR
     | Word -> Word
forall {a}. (Bits a, Num a) => a -> a
smallest Word
pw1 Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
pw2 -> TreeTreeBranch
ABL
     | Bool
otherwise           -> TreeTreeBranch
NOM
  Ordering
EQ                       -> TreeTreeBranch
EQL
  where
    pw1 :: Word
pw1 = Int -> Word
i2w (Prefix -> Int
unPrefix Prefix
p1)
    pw2 :: Word
pw2 = Int -> Word
i2w (Prefix -> Int
unPrefix Prefix
p2)
    greatest :: a -> a
greatest a
pw = a
pw a -> a -> a
forall a. Bits a => a -> a -> a
.|. (a
pwa -> a -> a
forall a. Num a => a -> a -> a
-a
1)
    smallest :: a -> a
smallest a
pw = a
pw a -> a -> a
forall a. Bits a => a -> a -> a
.&. (a
pwa -> a -> a
forall a. Num a => a -> a -> a
-a
1)
{-# INLINE treeTreeBranch #-}

-- | Whether this @Prefix@ splits a @Bin@ at the sign bit.
--
-- This can only be True at the top level.
-- If it is true, the left child contains non-negative keys and the right child
-- contains negative keys.
signBranch :: Prefix -> Bool
signBranch :: Prefix -> Bool
signBranch Prefix
p = Prefix -> Int
unPrefix Prefix
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
forall a. Bounded a => a
minBound :: Int)
{-# INLINE signBranch #-}

-- | The prefix of key @i@ up to (but not including) the switching
-- bit @m@.
mask :: Key -> Int -> Int
mask :: Int -> Int -> Int
mask Int
i Int
m = Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. ((-Int
m) Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
m)
{-# INLINE mask #-}

-- | The first switching bit where the two prefixes disagree.
--
-- Precondition for defined behavior: p1 /= p2
branchMask :: Int -> Int -> Int
branchMask :: Int -> Int -> Int
branchMask Int
p1 Int
p2 =
  Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
1 (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int
p1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
p2))
{-# INLINE branchMask #-}

i2w :: Int -> Word
i2w :: Int -> Word
i2w = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE i2w #-}

-- Used to compare IntSets and IntMaps
data Order
  = A_LT_B     -- holds for [0,3,4] [0,3,5,1]
  | A_Prefix_B -- holds for [0,3,4] [0,3,4,5]
  | A_EQ_B     -- holds for [0,3,4] [0,3,4]
  | B_Prefix_A -- holds for [0,3,4] [0,3]
  | A_GT_B     -- holds for [0,3,4] [0,2,5]

{--------------------------------------------------------------------
  Notes
--------------------------------------------------------------------}

-- Note [INLINE bit fiddling]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- It is essential that the bit fiddling functions like nomatch, mask,
-- branchMask etc are inlined. If they do not, the memory allocation skyrockets.
-- The GHC usually gets it right, but it is disastrous if it does not. Therefore
-- we explicitly mark these functions INLINE.