RangeTrie.hs

{-# LANGUAGE BangPatterns    #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns, PatternSynonyms #-}
{-# LANGUAGE GeneralizedNewtypeDeriving, DerivingVia, DerivingStrategies #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}

import Data.Bits
import Data.Bool
import Test.QuickCheck hiding ((.&.))
import Control.Applicative

size :: Int
size = 32

bits :: Int -> [Bool]
bits i = map (testBit i) [size-1,size-2..0]

truncTo :: Int -> Int -> Int
n `truncTo` t = n .&. (2 ^ t - 1)

trunc :: Int -> Int
trunc = flip truncTo size

data Ranges
  = None
  | !Ranges :!!: !Ranges
  | All
  | Seg {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !Int !Ranges
  --                    ^                   ^                   ^
  --                    Length              Path                Offshoots
  deriving (Show)
  
-- A segment in a tree is a run of children with only one non-leaf child, i.e.:
--
--       .
--       |\
--       | \
--       . All
--      /|
--     / |
-- None  .
--       |\              ~= (None :!: ((... :!: None) :!: All)) :!: All
--       | \
--       .  All
--       |\
--       | \
--       .  None
--      / \
--      ...
--
-- Here, it's wasteful to store this as a tree, since it's really just a linked list.
-- We can compress it, however, into a few `Int`s, using the `Seg` constructor.
-- This has 4 fields:
--
--   1. The first is the length of the segment. In the example above, it's 4.
--   2. The second is the path the tree continues down. Above, that's
--           Left, Right, Left, Left
--      We can encode this as a bit string, 0100, or the int 4.
--   3. The last field is the offshoots: the leaves as you walk down the path.
--      Above, they are:
--           All, None, All, None
--      Which encoded is 1010, giving us 10.
--
-- This means that instead of 4 constructors we can instead have the single constructor
--
--     Seg 4 4 10 ...

instance Eq Ranges where
  None == None = True
  All  == All  = True
  (xl :!: xr) == (yl :!: yr) = (xl == yl) && (xr == yr)
  _ == _ = False
  
instance Ord Ranges where
  None        <= _    = True
  _           <= All  = True
  All         <= _    = False
  _           <= None = False
  (xl :!: xr) <= (yl :!: yr) = case compare xl yl of
    LT -> True
    GT -> False
    EQ -> xr <= yr
    
  compare None None = EQ
  compare None _    = LT
  compare _    None = GT
  compare All  All  = EQ
  compare All  _    = GT
  compare _    All  = LT
  compare (xl :!: xr) (yl :!: yr) = compare xl yl <> compare xr yr

pattern (:!:) :: Ranges -> Ranges -> Ranges
pattern x :!: y <- (unconsBranch -> Just (x, y))
  where
    All         :!: All         = All
    None        :!: None        = None
    None        :!: Seg l p o r = Seg (l+1) (setBit p l) o            r
    All         :!: Seg l p o r = Seg (l+1) (setBit p l) (setBit o l) r
    Seg l p o r :!: None        = Seg (l+1) p            o            r
    Seg l p o r :!: All         = Seg (l+1) p            (setBit o l) r
    None        :!: y           = Seg 1 1 0 y
    All         :!: y           = Seg 1 1 1 y
    x           :!: None        = Seg 1 0 0 x
    x           :!: All         = Seg 1 0 1 x
    x           :!: y           = x :!!: y
{-# COMPLETE (:!:), None, All #-}

unconsBranch :: Ranges -> Maybe (Ranges, Ranges)
unconsBranch (xs :!!: ys) = Just (xs, ys)
unconsBranch (Seg l p o r) 
  | testBit p (l-1) = Just (bool None All (testBit o (l-1)), tl l p o r)
  | otherwise       = Just (tl l p o r, bool None All (testBit o (l-1)))
  where
    tl 1 _ _ r = r
    tl l p o r = Seg (l-1) (clearBit p (l-1)) (clearBit o (l-1)) r
unconsBranch None = Nothing
unconsBranch All = Nothing

compl :: Ranges -> Ranges
compl None = All
compl All  = None
compl (x :!: y) = compl x :!: compl y

singleton :: Int -> Ranges
singleton i = Seg size (trunc i) 0 All

differAt :: Int -> Int -> Int
differAt x y = (finiteBitSize x - 1) - countLeadingZeros (xor x y)

(?) :: Int -> Ranges -> Bool
(?) n = go (size-1)
  where
    go _ All  = True
    go _ None = False
    go i (l :!!: r)
      | testBit n i = go (i-1) r
      | otherwise   = go (i-1) l
      
    -- This finds if the needle diverges from the path, and if so it returns the
    -- offshoot at that point, otherwise (i.e. the path is followed to the end) we
    -- recurse on the child node.
    go i (Seg l p o r)
      | d >= 0 = testBit o d 
      | otherwise = go (i-l) r
      where d = differAt ((n `shiftR` ((i+1)-l)) `truncTo` l) p

slowMember :: Int -> Ranges -> Bool
slowMember i t = foldr f (All ==) (bits (trunc i)) t
  where
    f _     _ All       = True
    f _     _ None      = False
    f False k (l :!: _) = k l
    f True  k (_ :!: r) = k r

instance Num Ranges where
  None      + x         = x
  All       + _         = All
  x         + None      = x
  _         + All       = All
  xl :!: xr + yl :!: yr = (xl + yl) :!: (xr + yr)
  
  None      * _         = None
  All       * x         = x
  _         * None      = None
  x         * All       = x
  xl :!: xr * yl :!: yr = (xl * yl) :!: (xr * yr)
  
  fromInteger = singleton . fromEnum
  
  abs = id
  
  _         - All       = None
  x         - None      = x
  None      - _         = None
  All       - y         = compl y
  xl :!: xr - yl :!: yr = (xl - yl) :!: (xr - yr)
  
  signum None = 0
  signum _    = 1

atLeast :: Int -> Ranges
atLeast i = foldr (bool (:!: All) (None :!:)) All (take (size - countTrailingZeros i) (bits i))

lessThan :: Int -> Ranges
lessThan i = compl (atLeast i)

atMost :: Int -> Ranges
atMost i = lessThan i + singleton i

range :: Int -> Int -> Ranges
range lb ub = atLeast lb * atMost ub
    
enumerate :: Ranges -> [Int]
enumerate r = go size 0 r []
  where
    go !n !a All  = (++) [a .. (a + (2^n) - 1)]
    go !_ !_ None = id
    go !0 !_ _    = id
    go !n !a (l :!: r) = go (n-1) a l . go (n-1) ((2^(n-1)) + a) r

instance Arbitrary Ranges where
  arbitrary = sized (go . min size . fromEnum . logBase 2 . (toEnum :: Int -> Double))
    where
      go 0 = elements [All, None]
      go n = frequency [(n, let r = go (n-1) in liftA2 (:!:) r r), (1, elements [All, None])]
      
  shrink (x :!: y) = None : All : x : y : map (uncurry (:!:)) (shrink (x, y))
  shrink _ = []

newtype InRange = InRange Int 
  deriving stock (Eq, Ord) 
  deriving newtype (Show, Enum, Num, Integral, Real)

instance Bounded InRange where
  minBound = 0
  maxBound = 2 ^ size - 1
  
instance Arbitrary InRange where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink = shrinkIntegral

prop_member :: InRange -> Property
prop_member (InRange i) = property (i ? singleton i)

prop_fastMember :: InRange -> Ranges -> Property
prop_fastMember (InRange i) r = slowMember i r === (i ? r)

prop_delete :: InRange -> Ranges -> Property
prop_delete (InRange i) r = not (i ? r) ==> ((singleton i + r) - singleton i) === r

prop_compl :: Ranges -> Property
prop_compl r = r + compl r === All

prop_range :: InRange -> InRange -> InRange -> Property
prop_range (InRange lb) (InRange ub) (InRange i) = ((lb <= i) && (i  <= ub)) === (i ? range lb ub)

prop_all :: Property
prop_all = range 0 (2^size - 1) === All

return []

main :: IO Bool
main = $quickCheckAll