{-# LANGUAGE ConstrainedClassMethods #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

module Data.Relation
  ( Relation
      ( (⨃),
        (∪),
        dom,
        range,
        (◁),
        (<|),
        (▷),
        (|>),
        singleton,
        (⋪),
        (</|),
        (⋫),
        (|/>),
        Domain,
        Range,
        haskey,
        addpair,
        removekey,
        -- below are methods not used anywhere
        size
      ),
    (⊆),
    (∪+),
    (∈),
    (∉),
    (∩),
  )
where

import Data.Foldable (toList)
import Data.Kind (Type)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set, intersection, isSubsetOf)
import qualified Data.Set as Set

---------------------------------------------------------------------------------
-- Domain restriction and exclusion
---------------------------------------------------------------------------------

class Relation m where
  type Domain m :: Type
  type Range m :: Type

  singleton :: Domain m -> Range m -> m

  -- | Domain
  dom :: Ord (Domain m) => m -> Set (Domain m)

  -- | Range
  range :: Ord (Range m) => m -> Set (Range m)

  -- | Domain restriction
  --
  -- Unicode: 25c1
  (◁), (<|) :: (Ord (Domain m)) => Set (Domain m) -> m -> m
  Set (Domain m)
s <| m
r = Set (Domain m)
s Set (Domain m) -> m -> m
forall m. (Relation m, Ord (Domain m)) => Set (Domain m) -> m -> m
 m
r

  -- | Domain exclusion
  --
  -- Unicode: 22ea
  (⋪), (</|) :: (Ord (Domain m)) => Set (Domain m) -> m -> m
  Set (Domain m)
s </| m
r = Set (Domain m)
s Set (Domain m) -> m -> m
forall m. (Relation m, Ord (Domain m)) => Set (Domain m) -> m -> m
 m
r

  -- | Range restriction
  --
  -- Unicode: 25b7
  (▷), (|>) :: Ord (Range m) => m -> Set (Range m) -> m
  m
s |> Set (Range m)
r = m
s m -> Set (Range m) -> m
forall m. (Relation m, Ord (Range m)) => m -> Set (Range m) -> m
 Set (Range m)
r

  -- | Range exclusion
  --
  -- Unicode: 22eb
  (⋫), (|/>) :: Ord (Range m) => m -> Set (Range m) -> m
  m
s |/> Set (Range m)
r = m
s m -> Set (Range m) -> m
forall m. (Relation m, Ord (Range m)) => m -> Set (Range m) -> m
 Set (Range m)
r

  -- | Union
  (∪) :: (Ord (Domain m), Ord (Range m)) => m -> m -> m

  -- | Union Override Right
  (⨃) :: (Ord (Domain m), Ord (Range m)) => m -> m -> m

  -- | Size of the relation
  size :: Integral n => m -> n

  -- | Is this key in the Domain,  Instances should overide this default with
  -- something more efficient
  haskey :: Ord (Domain m) => Domain m -> m -> Bool
  haskey Domain m
key m
m = Domain m
key Domain m -> Set (Domain m) -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (m -> Set (Domain m)
forall m. (Relation m, Ord (Domain m)) => m -> Set (Domain m)
dom m
m)

  -- | Insert (key,value) pair into the Relation.  Instances should overide this
  -- default with something more efficient
  addpair :: (Ord (Domain m), Ord (Range m)) => Domain m -> Range m -> m -> m
  addpair Domain m
key Range m
val m
m = m
m m -> m -> m
forall m.
(Relation m, Ord (Domain m), Ord (Range m)) =>
m -> m -> m
 (Domain m -> Range m -> m
forall m. Relation m => Domain m -> Range m -> m
singleton Domain m
key Range m
val)

  -- | Remove a key (and its associted value at that key) from the Relation.
  -- Instances should overide this default with something more efficient
  removekey :: Ord (Domain m) => Domain m -> m -> m
  removekey Domain m
k m
m = Domain m -> Set (Domain m)
forall a. a -> Set a
Set.singleton Domain m
k Set (Domain m) -> m -> m
forall m. (Relation m, Ord (Domain m)) => Set (Domain m) -> m -> m
 m
m

-- | Alias for 'elem'.
--
-- Unicode: 2208
(∈) :: (Eq a, Foldable f) => a -> f a -> Bool
a
a ∈ :: a -> f a -> Bool
 f a
f = a -> f a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem a
a f a
f

-- | Alias for not 'elem'.
--
-- Unicode: 2209
(∉) :: (Eq a, Foldable f) => a -> f a -> Bool
a
a ∉ :: a -> f a -> Bool
 f a
f = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a -> f a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem a
a f a
f

infixl 4 

instance Relation (Map k v) where
  type Domain (Map k v) = k
  type Range (Map k v) = v

  singleton :: Domain (Map k v) -> Range (Map k v) -> Map k v
singleton = Domain (Map k v) -> Range (Map k v) -> Map k v
forall k a. k -> a -> Map k a
Map.singleton

  dom :: Map k v -> Set (Domain (Map k v))
dom = Map k v -> Set (Domain (Map k v))
forall k a. Map k a -> Set k
Map.keysSet
  range :: Map k v -> Set (Range (Map k v))
range = [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList ([v] -> Set v) -> (Map k v -> [v]) -> Map k v -> Set v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> [v]
forall k a. Map k a -> [a]
Map.elems

  Set (Domain (Map k v))
s ◁ :: Set (Domain (Map k v)) -> Map k v -> Map k v
 Map k v
r = Map k v -> Set k -> Map k v
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.restrictKeys Map k v
r Set k
Set (Domain (Map k v))
s

  Set (Domain (Map k v))
s ⋪ :: Set (Domain (Map k v)) -> Map k v -> Map k v
 Map k v
r = Map k v -> Set k -> Map k v
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map k v
r Set k
Set (Domain (Map k v))
s -- Uses library fuction which is equivalent to: Map.filterWithKey (\k _ -> k `Set.notMember` s) r

  Map k v
r ▷ :: Map k v -> Set (Range (Map k v)) -> Map k v
 Set (Range (Map k v))
s = (v -> Bool) -> Map k v -> Map k v
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set v
Set (Range (Map k v))
s) Map k v
r

  Map k v
r ⋫ :: Map k v -> Set (Range (Map k v)) -> Map k v
 Set (Range (Map k v))
s = (v -> Bool) -> Map k v -> Map k v
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set v
Set (Range (Map k v))
s) Map k v
r

  Map k v
d0 ∪ :: Map k v -> Map k v -> Map k v
 Map k v
d1 = Map k v -> Map k v -> Map k v
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k v
d0 Map k v
d1

  -- For union override we pass @d1@ as first argument, since 'Map.union' is left biased.
  Map k v
d0 ⨃ :: Map k v -> Map k v -> Map k v
 Map k v
d1 = Map k v -> Map k v -> Map k v
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k v
d1 Map k v
d0

  size :: Map k v -> n
size = Int -> n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> n) -> (Map k v -> Int) -> Map k v -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> Int
forall k a. Map k a -> Int
Map.size

  {-# INLINE haskey #-}
  haskey :: Domain (Map k v) -> Map k v -> Bool
haskey Domain (Map k v)
x Map k v
m = case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
Domain (Map k v)
x Map k v
m of Just v
_ -> Bool
True; Maybe v
Nothing -> Bool
False

  {-# INLINE addpair #-}
  addpair :: Domain (Map k v) -> Range (Map k v) -> Map k v -> Map k v
addpair = (v -> v -> v) -> k -> v -> Map k v -> Map k v
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith (\v
x v
_y -> v
x)

  {-# INLINE removekey #-}
  removekey :: Domain (Map k v) -> Map k v -> Map k v
removekey Domain (Map k v)
k Map k v
m = k -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
Domain (Map k v)
k Map k v
m

-- | Union override plus is (A\B)∪(B\A)∪{k|->v1+v2 | k|->v1 : A /\ k|->v2 : B}
-- The library function Map.unionWith is more general, it allows any type for
-- `b` as long as (+) :: b -> b -> b
(∪+) :: (Ord a, Num b) => Map a b -> Map a b -> Map a b
Map a b
a ∪+ :: Map a b -> Map a b -> Map a b
∪+ Map a b
b = ((b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith b -> b -> b
forall a. Num a => a -> a -> a
(+) Map a b
a Map a b
b)

instance Relation (Set (a, b)) where
  type Domain (Set (a, b)) = a
  type Range (Set (a, b)) = b

  singleton :: Domain (Set (a, b)) -> Range (Set (a, b)) -> Set (a, b)
singleton Domain (Set (a, b))
a Range (Set (a, b))
b = (a, b) -> Set (a, b)
forall a. a -> Set a
Set.singleton (a
Domain (Set (a, b))
a, b
Range (Set (a, b))
b)

  dom :: Set (a, b) -> Set (Domain (Set (a, b)))
dom = ((a, b) -> a) -> Set (a, b) -> Set a
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (a, b) -> a
forall a b. (a, b) -> a
fst

  range :: Set (a, b) -> Set (Range (Set (a, b)))
range = ((a, b) -> b) -> Set (a, b) -> Set b
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (a, b) -> b
forall a b. (a, b) -> b
snd

  Set (Domain (Set (a, b)))
s ◁ :: Set (Domain (Set (a, b))) -> Set (a, b) -> Set (a, b)
 Set (a, b)
r = ((a, b) -> Bool) -> Set (a, b) -> Set (a, b)
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\(a
k, b
_) -> a
k a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set a
Set (Domain (Set (a, b)))
s) Set (a, b)
r

  Set (Domain (Set (a, b)))
s ⋪ :: Set (Domain (Set (a, b))) -> Set (a, b) -> Set (a, b)
 Set (a, b)
r = ((a, b) -> Bool) -> Set (a, b) -> Set (a, b)
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\(a
k, b
_) -> a
k a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set a
Set (Domain (Set (a, b)))
s) Set (a, b)
r

  Set (a, b)
r ▷ :: Set (a, b) -> Set (Range (Set (a, b))) -> Set (a, b)
 Set (Range (Set (a, b)))
s = ((a, b) -> Bool) -> Set (a, b) -> Set (a, b)
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\(a
_, b
v) -> b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member b
v Set b
Set (Range (Set (a, b)))
s) Set (a, b)
r

  Set (a, b)
r ⋫ :: Set (a, b) -> Set (Range (Set (a, b))) -> Set (a, b)
 Set (Range (Set (a, b)))
s = ((a, b) -> Bool) -> Set (a, b) -> Set (a, b)
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\(a
_, b
v) -> b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.notMember b
v Set b
Set (Range (Set (a, b)))
s) Set (a, b)
r

  ∪ :: Set (a, b) -> Set (a, b) -> Set (a, b)
(∪) = Set (a, b) -> Set (a, b) -> Set (a, b)
forall a. Ord a => Set a -> Set a -> Set a
Set.union

  Set (a, b)
d0 ⨃ :: Set (a, b) -> Set (a, b) -> Set (a, b)
 Set (a, b)
d1 = Set (a, b)
d1' Set (a, b) -> Set (a, b) -> Set (a, b)
forall m.
(Relation m, Ord (Domain m), Ord (Range m)) =>
m -> m -> m
 ((Set (a, b) -> Set (Domain (Set (a, b)))
forall m. (Relation m, Ord (Domain m)) => m -> Set (Domain m)
dom Set (a, b)
d1') Set (Domain (Set (a, b))) -> Set (a, b) -> Set (a, b)
forall m. (Relation m, Ord (Domain m)) => Set (Domain m) -> m -> m
 Set (a, b)
d0)
    where
      d1' :: Set (a, b)
d1' = Set (a, b) -> Set (a, b)
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set (a, b)
d1

  size :: Set (a, b) -> n
size = Int -> n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> n) -> (Set (a, b) -> Int) -> Set (a, b) -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (a, b) -> Int
forall a. Set a -> Int
Set.size

  addpair :: Domain (Set (a, b))
-> Range (Set (a, b)) -> Set (a, b) -> Set (a, b)
addpair Domain (Set (a, b))
key Range (Set (a, b))
val Set (a, b)
set = (a, b) -> Set (a, b) -> Set (a, b)
forall a. Ord a => a -> Set a -> Set a
Set.insert (a
Domain (Set (a, b))
key, b
Range (Set (a, b))
val) Set (a, b)
set

-- The [(a,b)] instance is used in `stakeDistr` in the file LedgerState.hs
instance Relation [(a, b)] where
  type Domain [(a, b)] = a
  type Range [(a, b)] = b

  singleton :: Domain [(a, b)] -> Range [(a, b)] -> [(a, b)]
singleton Domain [(a, b)]
a Range [(a, b)]
b = [(a
Domain [(a, b)]
a, b
Range [(a, b)]
b)]

  dom :: [(a, b)] -> Set (Domain [(a, b)])
dom = [a] -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet ([a] -> Set a) -> ([(a, b)] -> [a]) -> [(a, b)] -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> a) -> [(a, b)] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, b) -> a
forall a b. (a, b) -> a
fst

  range :: [(a, b)] -> Set (Range [(a, b)])
range = [b] -> Set b
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet ([b] -> Set b) -> ([(a, b)] -> [b]) -> [(a, b)] -> Set b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> b) -> [(a, b)] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, b) -> b
forall a b. (a, b) -> b
snd

  Set (Domain [(a, b)])
s ◁ :: Set (Domain [(a, b)]) -> [(a, b)] -> [(a, b)]
 [(a, b)]
r = ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set a
Set (Domain [(a, b)])
s) (a -> Bool) -> ((a, b) -> a) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst) [(a, b)]
r

  Set (Domain [(a, b)])
s ⋪ :: Set (Domain [(a, b)]) -> [(a, b)] -> [(a, b)]
 [(a, b)]
r = ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set a
Set (Domain [(a, b)])
s) (a -> Bool) -> ((a, b) -> a) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst) [(a, b)]
r

  [(a, b)]
r ▷ :: [(a, b)] -> Set (Range [(a, b)]) -> [(a, b)]
 Set (Range [(a, b)])
s = ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set b -> Set b
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set b
Set (Range [(a, b)])
s) (b -> Bool) -> ((a, b) -> b) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> b
forall a b. (a, b) -> b
snd) [(a, b)]
r

  [(a, b)]
r ⋫ :: [(a, b)] -> Set (Range [(a, b)]) -> [(a, b)]
 Set (Range [(a, b)])
s = ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set b -> Set b
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet Set b
Set (Range [(a, b)])
s) (b -> Bool) -> ((a, b) -> b) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> b
forall a b. (a, b) -> b
snd) [(a, b)]
r

  ∪ :: [(a, b)] -> [(a, b)] -> [(a, b)]
(∪) = [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. [a] -> [a] -> [a]
(++)

  -- In principle a list of pairs allows for duplicated keys.
  [(a, b)]
d0 ⨃ :: [(a, b)] -> [(a, b)] -> [(a, b)]
 [(a, b)]
d1 = [(a, b)]
d0 [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. [a] -> [a] -> [a]
++ [(a, b)] -> [(a, b)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [(a, b)]
d1

  size :: [(a, b)] -> n
size = Int -> n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> n) -> ([(a, b)] -> Int) -> [(a, b)] -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, b)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

  addpair :: Domain [(a, b)] -> Range [(a, b)] -> [(a, b)] -> [(a, b)]
addpair Domain [(a, b)]
key Range [(a, b)]
val [(a, b)]
list = (a
Domain [(a, b)]
key, b
Range [(a, b)]
val) (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: [(a, b)]
list

---------------------------------------------------------------------------------
-- Aliases
---------------------------------------------------------------------------------

-- | Inclusion among foldables.
--
-- Unicode: 2286
(⊆) :: (Foldable f, Foldable g, Ord a) => f a -> g a -> Bool
f a
x ⊆ :: f a -> g a -> Bool
 g a
y = f a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet f a
x Set a -> Set a -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`isSubsetOf` g a -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Set a
toSet g a
y

toSet :: (Foldable f, Ord a) => f a -> Set a
toSet :: f a -> Set a
toSet = [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> (f a -> [a]) -> f a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

(∩) :: Ord a => Set a -> Set a -> Set a
∩ :: Set a -> Set a -> Set a
(∩) = Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
intersection