{-# LANGUAGE Safe #-}
{-# OPTIONS_HADDOCK not-home #-}
module Kleene.Internal.Partition where

import Data.Foldable      (toList)
import Data.List.NonEmpty (NonEmpty (..))
import Data.RangeSet.Map  (RSet)
import Data.Set           (Set)

import qualified Data.Function.Step.Discrete.Closed as SF
import qualified Data.List.NonEmpty                 as NE
import qualified Data.RangeSet.Map                  as RSet
import qualified Data.Set                           as Set

import Test.QuickCheck

-- $setup
-- >>> import Data.Word
-- >>> import Data.Semigroup (Semigroup (..))
-- >>> import Test.QuickCheck (Arbitrary (..), (===), (==>))
-- >>> import Data.RangeSet.Map (RSet)
-- >>> import qualified Data.RangeSet.Map as RSet
--
-- >>> let asPartitionChar :: Partition Char -> Partition Char; asPartitionChar = id
-- >>> instance (Ord a, Enum a, Arbitrary a) => Arbitrary (RSet a) where arbitrary = fmap RSet.fromRangeList arbitrary

-- | 'Partition' devides type into disjoint connected partitions.
--
-- /Note:/ we could have non-connecter partitions too,
-- but that would be more complicated.
-- This variant is correct by construction, but less precise.
--
-- It's enought to store last element of each piece.
--
-- @'Partition' (fromList [x1, x2, x3]) :: 'Partition' s@ describes a partition of /Set/ @s@, as
--
-- \[
-- \{ x \mid x \le x_1 \} \cup
-- \{ x \mid x_1 < x \le x_2 \} \cup
-- \{ x \mid x_2 < x \le x_3 \} \cup
-- \{ x \mid x_3 < x \}
-- \]
--
-- /Note:/ it's enough to check upper bound conditions only if checks are performed in order.
--
-- /Invariant:/ 'maxBound' is not in the set.
--
newtype Partition a  = Partition { forall a. Partition a -> Set a
unPartition :: Set a }
  deriving (Partition a -> Partition a -> Bool
(Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool) -> Eq (Partition a)
forall a. Eq a => Partition a -> Partition a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Partition a -> Partition a -> Bool
== :: Partition a -> Partition a -> Bool
$c/= :: forall a. Eq a => Partition a -> Partition a -> Bool
/= :: Partition a -> Partition a -> Bool
Eq, Eq (Partition a)
Eq (Partition a) =>
(Partition a -> Partition a -> Ordering)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Partition a)
-> (Partition a -> Partition a -> Partition a)
-> Ord (Partition a)
Partition a -> Partition a -> Bool
Partition a -> Partition a -> Ordering
Partition a -> Partition a -> Partition a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Partition a)
forall a. Ord a => Partition a -> Partition a -> Bool
forall a. Ord a => Partition a -> Partition a -> Ordering
forall a. Ord a => Partition a -> Partition a -> Partition a
$ccompare :: forall a. Ord a => Partition a -> Partition a -> Ordering
compare :: Partition a -> Partition a -> Ordering
$c< :: forall a. Ord a => Partition a -> Partition a -> Bool
< :: Partition a -> Partition a -> Bool
$c<= :: forall a. Ord a => Partition a -> Partition a -> Bool
<= :: Partition a -> Partition a -> Bool
$c> :: forall a. Ord a => Partition a -> Partition a -> Bool
> :: Partition a -> Partition a -> Bool
$c>= :: forall a. Ord a => Partition a -> Partition a -> Bool
>= :: Partition a -> Partition a -> Bool
$cmax :: forall a. Ord a => Partition a -> Partition a -> Partition a
max :: Partition a -> Partition a -> Partition a
$cmin :: forall a. Ord a => Partition a -> Partition a -> Partition a
min :: Partition a -> Partition a -> Partition a
Ord)

-- | Check invariant.
invariant :: (Ord a, Bounded a) => Partition a -> Bool
invariant :: forall a. (Ord a, Bounded a) => Partition a -> Bool
invariant (Partition Set a
xs) = a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.notMember a
forall a. Bounded a => a
maxBound Set a
xs

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

instance Show a => Show (Partition a) where
    showsPrec :: Int -> Partition a -> ShowS
showsPrec Int
d (Partition Set a
xs)
        = Bool -> ShowS -> ShowS
showParen (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10)
        (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"fromSeparators "
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
xs)

-- | prop> invariant (asPartitionChar p)
instance (Enum a, Bounded a, Ord a, Arbitrary a) => Arbitrary (Partition a) where
    arbitrary :: Gen (Partition a)
arbitrary = [a] -> Partition a
forall a. (Enum a, Bounded a, Ord a) => [a] -> Partition a
fromSeparators ([a] -> Partition a) -> Gen [a] -> Gen (Partition a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [a]
forall a. Arbitrary a => Gen a
arbitrary

-- | See 'wedge'.
instance (Enum a, Bounded a, Ord a) => Semigroup (Partition a) where
    <> :: Partition a -> Partition a -> Partition a
(<>) = Partition a -> Partition a -> Partition a
forall a. Ord a => Partition a -> Partition a -> Partition a
wedge

instance (Enum a, Bounded a, Ord a) => Monoid (Partition a) where
    mempty :: Partition a
mempty = Partition a
forall a. Partition a
whole
    mappend :: Partition a -> Partition a -> Partition a
mappend = Partition a -> Partition a -> Partition a
forall a. Semigroup a => a -> a -> a
(<>)

-------------------------------------------------------------------------------
-- Constructors
-------------------------------------------------------------------------------

fromSeparators :: (Enum a, Bounded a, Ord a) => [a] -> Partition a
fromSeparators :: forall a. (Enum a, Bounded a, Ord a) => [a] -> Partition a
fromSeparators = Set a -> Partition a
forall a. Set a -> Partition a
Partition (Set a -> Partition a) -> ([a] -> Set a) -> [a] -> Partition a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> ([a] -> [a]) -> [a] -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
forall a. Bounded a => a
maxBound)

-- | Construct 'Partition' from list of 'RSet's.
--
-- RSet intervals are closed on both sides.
fromRSets :: (Enum a, Bounded a, Ord a) => [RSet a] -> Partition a
fromRSets :: forall a. (Enum a, Bounded a, Ord a) => [RSet a] -> Partition a
fromRSets [RSet a]
rs = Set a -> Partition a
forall a. Set a -> Partition a
Partition (Set a -> Partition a) -> Set a -> Partition a
forall a b. (a -> b) -> a -> b
$ [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ (if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Bounded a => a
minBound then [] else [a -> a
forall a. Enum a => a -> a
pred a
x]) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
      (if a
y a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Bounded a => a
maxBound then [] else [a
y])
    | RSet a
r <- [RSet a]
rs
    , (a
x, a
y) <- RSet a -> [(a, a)]
forall a. RSet a -> [(a, a)]
RSet.toRangeList RSet a
r
    ]

fromRSet :: (Enum a, Bounded a, Ord a) => RSet a -> Partition a
fromRSet :: forall a. (Enum a, Bounded a, Ord a) => RSet a -> Partition a
fromRSet RSet a
r
    | RSet a
r RSet a -> RSet a -> Bool
forall a. Eq a => a -> a -> Bool
== RSet a
forall a. RSet a
RSet.empty = Partition a
forall a. Partition a
whole
    | RSet a
r RSet a -> RSet a -> Bool
forall a. Eq a => a -> a -> Bool
== RSet a
forall a. Bounded a => RSet a
RSet.full  = Partition a
forall a. Partition a
whole
    | Bool
otherwise       = [RSet a] -> Partition a
forall a. (Enum a, Bounded a, Ord a) => [RSet a] -> Partition a
fromRSets [RSet a
r]

whole :: Partition a
whole :: forall a. Partition a
whole = Set a -> Partition a
forall a. Set a -> Partition a
Partition Set a
forall a. Set a
Set.empty

-------------------------------------------------------------------------------
-- Querying
-------------------------------------------------------------------------------

-- | Count of sets in a 'Partition'.
--
-- >>> size whole
-- 1
--
-- >>> size $ split (10 :: Word8)
-- 2
--
-- prop> size (asPartitionChar p) >= 1
--
size :: Partition a -> Int
size :: forall a. Partition a -> Int
size (Partition Set a
xs) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set a
xs

-- | Extract examples from each subset in a 'Partition'.
--
-- >>> examples $ split (10 :: Word8)
-- fromList [10,255]
--
-- >>> examples $ split (10 :: Word8) <> split 20
-- fromList [10,20,255]
--
-- prop> invariant p ==> size (asPartitionChar p) === length (examples p)
--
examples :: (Bounded a, Enum a, Ord a) => Partition a -> Set a
examples :: forall a. (Bounded a, Enum a, Ord a) => Partition a -> Set a
examples (Partition Set a
xs) = a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
forall a. Bounded a => a
maxBound Set a
xs

-- |
--
-- prop> all (uncurry (<=)) $ intervals $ asPartitionChar p
intervals :: (Enum a, Bounded a, Ord a) => Partition a -> NonEmpty (a, a)
intervals :: forall a.
(Enum a, Bounded a, Ord a) =>
Partition a -> NonEmpty (a, a)
intervals (Partition Set a
xs) = a -> [a] -> NonEmpty (a, a)
forall {t}. Bounded t => t -> [t] -> NonEmpty (t, t)
go a
forall a. Bounded a => a
minBound (Set a -> [a]
forall a. Set a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Set a
xs) where
    go :: t -> [t] -> NonEmpty (t, t)
go t
x []       = (t
x, t
forall a. Bounded a => a
maxBound) (t, t) -> [(t, t)] -> NonEmpty (t, t)
forall a. a -> [a] -> NonEmpty a
:| []
    go t
x (t
y : [t]
ys) = (t
x, t
y) (t, t) -> NonEmpty (t, t) -> NonEmpty (t, t)
forall a. a -> NonEmpty a -> NonEmpty a
`NE.cons` t -> [t] -> NonEmpty (t, t)
go t
y [t]
ys

-------------------------------------------------------------------------------
--
-- Operations
-------------------------------------------------------------------------------

-- | Wedge partitions.
--
-- >>> split (10 :: Word8) <> split 20
-- fromSeparators [10,20]
--
-- prop> whole `wedge` (p :: Partition Char) === p
-- prop> (p :: Partition Char) <> whole === p
-- prop> asPartitionChar p <> q === q <> p
-- prop> asPartitionChar p <> p === p
-- prop> invariant $ asPartitionChar p <> q
--
wedge :: Ord a => Partition a -> Partition a -> Partition a
wedge :: forall a. Ord a => Partition a -> Partition a -> Partition a
wedge (Partition Set a
as) (Partition Set a
bs) = Set a -> Partition a
forall a. Set a -> Partition a
Partition (Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
as Set a
bs)

-- | Simplest partition: given @x@ partition space into @[min..x) and [x .. max]@
--
-- >>> split (128 :: Word8)
-- fromSeparators [128]
--
split :: (Enum a, Bounded a, Eq a) => a -> Partition a
split :: forall a. (Enum a, Bounded a, Eq a) => a -> Partition a
split a
x
    | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Bounded a => a
minBound = Set a -> Partition a
forall a. Set a -> Partition a
Partition Set a
forall a. Set a
Set.empty
    | Bool
otherwise     = Set a -> Partition a
forall a. Set a -> Partition a
Partition (a -> Set a
forall a. a -> Set a
Set.singleton a
x)

-------------------------------------------------------------------------------
-- Conversion
-------------------------------------------------------------------------------

-- | Make a step function.
toSF :: (Enum a, Bounded a, Ord a) => (a -> b) -> Partition a -> SF.SF a b
toSF :: forall a b.
(Enum a, Bounded a, Ord a) =>
(a -> b) -> Partition a -> SF a b
toSF a -> b
f (Partition Set a
p) = [(a, b)] -> b -> SF a b
forall k v. Ord k => [(k, v)] -> v -> SF k v
SF.fromList
    ((a -> (a, b)) -> [a] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
k -> (a
k, a -> b
f a
k)) ([a] -> [(a, b)]) -> [a] -> [(a, b)]
forall a b. (a -> b) -> a -> b
$ [a] -> [a]
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [a]
as)
    (a -> b
f a
forall a. Bounded a => a
maxBound)
  where
    as :: [a]
as = Set a -> [a]
forall a. Set a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Set a
p