{-# LANGUAGE
      CPP,
      DerivingVia
  #-}

module Data.Mapping.Piecewise where

#if MIN_VERSION_GLASGOW_HASKELL(9,6,0,0)
#else
import Control.Applicative (liftA2)
#endif
import Control.Applicative (liftA3)
import Data.Algebra.Boolean
import qualified Data.Map.Internal as MI
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.Mapping


-- | A data structure storing mappings that are constant on
-- intervals.
--
-- If the space of keys not discrete, then these mappings are
-- right-continuous: values are in general defined on intervals $a
-- \leq x < b$ which are closed on the left and open on the right.
data Piecewise k v = Piecewise {
  -- | The value taken for sufficiently small keys
  forall k v. Piecewise k v -> v
leftEnd :: v,
  forall k v. Piecewise k v -> Map k v
starts :: Map k v
} deriving (Piecewise k v -> Piecewise k v -> Bool
(Piecewise k v -> Piecewise k v -> Bool)
-> (Piecewise k v -> Piecewise k v -> Bool) -> Eq (Piecewise k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
$c== :: forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
== :: Piecewise k v -> Piecewise k v -> Bool
$c/= :: forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
/= :: Piecewise k v -> Piecewise k v -> Bool
Eq, Eq (Piecewise k v)
Eq (Piecewise k v) =>
(Piecewise k v -> Piecewise k v -> Ordering)
-> (Piecewise k v -> Piecewise k v -> Bool)
-> (Piecewise k v -> Piecewise k v -> Bool)
-> (Piecewise k v -> Piecewise k v -> Bool)
-> (Piecewise k v -> Piecewise k v -> Bool)
-> (Piecewise k v -> Piecewise k v -> Piecewise k v)
-> (Piecewise k v -> Piecewise k v -> Piecewise k v)
-> Ord (Piecewise k v)
Piecewise k v -> Piecewise k v -> Bool
Piecewise k v -> Piecewise k v -> Ordering
Piecewise k v -> Piecewise k v -> Piecewise k v
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 k v. (Ord v, Ord k) => Eq (Piecewise k v)
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Ordering
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
$ccompare :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Ordering
compare :: Piecewise k v -> Piecewise k v -> Ordering
$c< :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
< :: Piecewise k v -> Piecewise k v -> Bool
$c<= :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
<= :: Piecewise k v -> Piecewise k v -> Bool
$c> :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
> :: Piecewise k v -> Piecewise k v -> Bool
$c>= :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
>= :: Piecewise k v -> Piecewise k v -> Bool
$cmax :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
max :: Piecewise k v -> Piecewise k v -> Piecewise k v
$cmin :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
min :: Piecewise k v -> Piecewise k v -> Piecewise k v
Ord)

-- | The value taken for sufficiently large keys
rightEnd :: Piecewise k v -> v
rightEnd :: forall k v. Piecewise k v -> v
rightEnd (Piecewise v
a Map k v
m) = case Map k v -> Maybe (k, v)
forall k a. Map k a -> Maybe (k, a)
M.lookupMax Map k v
m of
  Maybe (k, v)
Nothing    -> v
a
  Just (k
_,v
b) -> v
b

-- | Assumes the keys are distinct and increasing (but consecutive
-- values may be the same, in which case the intervening keys are
-- removed)
fromAscList :: (Eq v) => v -> [(k,v)] -> Piecewise k v
fromAscList :: forall v k. Eq v => v -> [(k, v)] -> Piecewise k v
fromAscList = let
  inner :: b -> [(a, b)] -> [(a, b)]
inner b
_ [] = []
  inner b
a ((a
y,b
b):[(a, b)]
r)
    | b
a b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
b    = b -> [(a, b)] -> [(a, b)]
inner b
a [(a, b)]
r
    | Bool
otherwise = (a
y,b
b)(a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
:b -> [(a, b)] -> [(a, b)]
inner b
b [(a, b)]
r
  run :: v -> [(k, v)] -> Piecewise k v
run v
x = v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
x (Map k v -> Piecewise k v)
-> ([(k, v)] -> Map k v) -> [(k, v)] -> Piecewise k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, v)] -> Map k v
forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList ([(k, v)] -> Map k v)
-> ([(k, v)] -> [(k, v)]) -> [(k, v)] -> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> [(k, v)] -> [(k, v)]
forall {b} {a}. Eq b => b -> [(a, b)] -> [(a, b)]
inner v
x
  in v -> [(k, v)] -> Piecewise k v
forall v k. Eq v => v -> [(k, v)] -> Piecewise k v
run

instance (Show k, Show v) => Show (Piecewise k v) where
  showsPrec :: Int -> Piecewise k v -> ShowS
showsPrec Int
d (Piecewise v
k Map k v
m) =
    (String
"fromAscList " <>) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    Int -> v -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
d v
k ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (String
" " <>) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    [(k, v)] -> ShowS
forall a. Show a => [a] -> ShowS
showList (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k v
m)

-- | Assumes that the keys are distinct and increasing, and also that
-- consecutive values are distinct
fromAscListUnsafe :: v -> [(k,v)] -> Piecewise k v
fromAscListUnsafe :: forall v k. v -> [(k, v)] -> Piecewise k v
fromAscListUnsafe v
k = v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
k (Map k v -> Piecewise k v)
-> ([(k, v)] -> Map k v) -> [(k, v)] -> Piecewise k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, v)] -> Map k v
forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList

-- | Takes value `a` for keys less than `x` and `b` otherwise
changeAt :: v -> k -> v -> Piecewise k v
changeAt :: forall v k. v -> k -> v -> Piecewise k v
changeAt v
a k
x v
b = v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a (Map k v -> Piecewise k v) -> Map k v -> Piecewise k v
forall a b. (a -> b) -> a -> b
$ k -> v -> Map k v
forall k a. k -> a -> Map k a
M.singleton k
x v
b

-- | Is the value greater than or equal to `k`?
greaterThanOrEqual :: k -> Piecewise k Bool
greaterThanOrEqual :: forall k. k -> Piecewise k Bool
greaterThanOrEqual k
k = Bool -> k -> Bool -> Piecewise k Bool
forall v k. v -> k -> v -> Piecewise k v
changeAt Bool
False k
k Bool
True

-- | Is the value less than `k`?
lessThan :: k -> Piecewise k Bool
lessThan :: forall k. k -> Piecewise k Bool
lessThan k
k = Bool -> k -> Bool -> Piecewise k Bool
forall v k. v -> k -> v -> Piecewise k v
changeAt Bool
True k
k Bool
False

-- | Is the value greater than `k`? This is subject to the usual
-- concerns about `Enum` (it not to be used with floating-point
-- arithmetic, for example)
greaterThan :: Enum k => k -> Piecewise k Bool
greaterThan :: forall k. Enum k => k -> Piecewise k Bool
greaterThan = k -> Piecewise k Bool
forall k. k -> Piecewise k Bool
greaterThanOrEqual (k -> Piecewise k Bool) -> (k -> k) -> k -> Piecewise k Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> k
forall a. Enum a => a -> a
succ

-- | Is the value less than or equal to `k`? This is subject to the
-- usual concerns about `Enum` (it not to be used with floating-point
-- arithmetic, for example)
lessThanOrEqual :: Enum k => k -> Piecewise k Bool
lessThanOrEqual :: forall k. Enum k => k -> Piecewise k Bool
lessThanOrEqual = k -> Piecewise k Bool
forall k. k -> Piecewise k Bool
lessThan (k -> Piecewise k Bool) -> (k -> k) -> k -> Piecewise k Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> k
forall a. Enum a => a -> a
succ

-- | All values, in order of increasing key
values :: Piecewise k v -> [v]
values :: forall k v. Piecewise k v -> [v]
values (Piecewise v
x Map k v
m) = v
x v -> [v] -> [v]
forall a. a -> [a] -> [a]
: Map k v -> [v]
forall k a. Map k a -> [a]
M.elems Map k v
m

instance (Eq k) => Functor (Piecewise k) where
  fmap :: forall a b. (a -> b) -> Piecewise k a -> Piecewise k b
fmap a -> b
p (Piecewise a
a Map k a
f) = b -> [(k, b)] -> Piecewise k b
forall v k. v -> [(k, v)] -> Piecewise k v
fromAscListUnsafe (a -> b
p a
a) ((a -> b) -> (k, a) -> (k, b)
forall a b. (a -> b) -> (k, a) -> (k, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
p ((k, a) -> (k, b)) -> [(k, a)] -> [(k, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map k a
f)

instance Foldable (Piecewise k) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Piecewise k a -> m
foldMap a -> m
f (Piecewise a
a Map k a
m) = a -> m
f a
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (a -> m) -> Map k a -> m
forall m a. Monoid m => (a -> m) -> Map k a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Map k a
m

instance Ord k => Mapping k (Piecewise k) where

  cst :: forall v. v -> Piecewise k v
cst v
x = v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
x Map k v
forall k a. Map k a
M.empty

  act :: forall v. Piecewise k v -> k -> v
act (Piecewise v
a Map k v
f) k
x = case k -> Map k v -> Maybe (k, v)
forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupLE k
x Map k v
f of
    Maybe (k, v)
Nothing -> v
a
    Just (k
_,v
b) -> v
b

  isConst :: forall v. Ord v => Piecewise k v -> Maybe v
isConst (Piecewise v
a Map k v
f) = if Map k v -> Bool
forall k a. Map k a -> Bool
M.null Map k v
f then v -> Maybe v
forall a. a -> Maybe a
Just v
a else Maybe v
forall a. Maybe a
Nothing

  mmap :: forall v u. Ord v => (u -> v) -> Piecewise k u -> Piecewise k v
mmap = (u -> v) -> Piecewise k u -> Piecewise k v
forall a b. (a -> b) -> Piecewise k a -> Piecewise k b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap

  mtraverse :: forall (f :: * -> *) v u.
(Applicative f, Ord v) =>
(u -> f v) -> Piecewise k u -> f (Piecewise k v)
mtraverse u -> f v
p (Piecewise u
a Map k u
f) = (v -> [(k, v)] -> Piecewise k v)
-> f v -> f [(k, v)] -> f (Piecewise k v)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 v -> [(k, v)] -> Piecewise k v
forall v k. Eq v => v -> [(k, v)] -> Piecewise k v
fromAscList (u -> f v
p u
a) (((k, u) -> f (k, v)) -> [(k, u)] -> f [(k, v)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((u -> f v) -> (k, u) -> f (k, v)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (k, a) -> f (k, b)
traverse u -> f v
p) ([(k, u)] -> f [(k, v)]) -> [(k, u)] -> f [(k, v)]
forall a b. (a -> b) -> a -> b
$ Map k u -> [(k, u)]
forall k a. Map k a -> [(k, a)]
M.toList Map k u
f)

  merge :: forall w u v.
Ord w =>
(u -> v -> w) -> Piecewise k u -> Piecewise k v -> Piecewise k w
merge u -> v -> w
p = let

    inner :: u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b w
c r :: [(a, u)]
r@((a
x,u
a'):[(a, u)]
r') s :: [(a, v)]
s@((a
y,v
b'):[(a, v)]
s') = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
      Ordering
LT -> let
        c' :: w
c' = u -> v -> w
p u
a' v
b
        in if w
c' w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c [(a, u)]
r' [(a, v)]
s else (a
x,w
c')(a, w) -> [(a, w)] -> [(a, w)]
forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c' [(a, u)]
r' [(a, v)]
s
      Ordering
GT -> let
        c' :: w
c' = u -> v -> w
p u
a v
b'
        in if w
c' w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c [(a, u)]
r [(a, v)]
s' else (a
y,w
c')(a, w) -> [(a, w)] -> [(a, w)]
forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c' [(a, u)]
r [(a, v)]
s'
      Ordering
EQ -> let
        c' :: w
c' = u -> v -> w
p u
a' v
b'
        in if w
c' w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b' w
c [(a, u)]
r' [(a, v)]
s' else (a
x,w
c')(a, w) -> [(a, w)] -> [(a, w)]
forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b' w
c' [(a, u)]
r' [(a, v)]
s'
    inner u
a v
_ w
c [] ((a
y,v
b'):[(a, v)]
s') = let
      c' :: w
c' = u -> v -> w
p u
a v
b'
      in if w
c' w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c [] [(a, v)]
s' else (a
y,w
c')(a, w) -> [(a, w)] -> [(a, w)]
forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c' [] [(a, v)]
s'
    inner u
_ v
b w
c ((a
x,u
a'):[(a, u)]
r') [] = let
      c' :: w
c' = u -> v -> w
p u
a' v
b
      in if w
c' w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c [(a, u)]
r' [] else (a
x,w
c')(a, w) -> [(a, w)] -> [(a, w)]
forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c' [(a, u)]
r' []
    inner u
_ v
_ w
_ [] [] = []

    run :: Piecewise k u -> Piecewise k v -> Piecewise k w
run (Piecewise u
a Map k u
f) (Piecewise v
b Map k v
g) = let
      c :: w
c = u -> v -> w
p u
a v
b
      l :: [(k, w)]
l = u -> v -> w -> [(k, u)] -> [(k, v)] -> [(k, w)]
forall {a}.
Ord a =>
u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b w
c (Map k u -> [(k, u)]
forall k a. Map k a -> [(k, a)]
M.toList Map k u
f) (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k v
g)
      in w -> Map k w -> Piecewise k w
forall k v. v -> Map k v -> Piecewise k v
Piecewise w
c (Map k w -> Piecewise k w) -> Map k w -> Piecewise k w
forall a b. (a -> b) -> a -> b
$ [(k, w)] -> Map k w
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k, w)]
l

    in Piecewise k u -> Piecewise k v -> Piecewise k w
forall {k}.
Ord k =>
Piecewise k u -> Piecewise k v -> Piecewise k w
run

  mergeA :: forall (f :: * -> *) w u v.
(Applicative f, Ord w) =>
(u -> v -> f w)
-> Piecewise k u -> Piecewise k v -> f (Piecewise k w)
mergeA u -> v -> f w
p = let

    maybePrepend :: a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x b
u b
v [(a, b)]
l
      | b
u b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
v    = [(a, b)]
l
      | Bool
otherwise = (a
x,b
v)(a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
:[(a, b)]
l

    inner :: u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b f w
c r :: [(a, u)]
r@((a
x,u
a'):[(a, u)]
r') s :: [(a, v)]
s@((a
y,v
b'):[(a, v)]
s') = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
      Ordering
LT -> let
        c' :: f w
c' = u -> v -> f w
p u
a' v
b
        in (w -> w -> [(a, w)] -> [(a, w)])
-> f w -> f w -> f [(a, w)] -> f [(a, w)]
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (a -> w -> w -> [(a, w)] -> [(a, w)]
forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' (f [(a, w)] -> f [(a, w)]) -> f [(a, w)] -> f [(a, w)]
forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b f w
c' [(a, u)]
r' [(a, v)]
s
      Ordering
GT -> let
        c' :: f w
c' = u -> v -> f w
p u
a v
b'
        in (w -> w -> [(a, w)] -> [(a, w)])
-> f w -> f w -> f [(a, w)] -> f [(a, w)]
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (a -> w -> w -> [(a, w)] -> [(a, w)]
forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
y) f w
c f w
c' (f [(a, w)] -> f [(a, w)]) -> f [(a, w)] -> f [(a, w)]
forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b' f w
c' [(a, u)]
r [(a, v)]
s'
      Ordering
EQ -> let
        c' :: f w
c' = u -> v -> f w
p u
a' v
b'
        in (w -> w -> [(a, w)] -> [(a, w)])
-> f w -> f w -> f [(a, w)] -> f [(a, w)]
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (a -> w -> w -> [(a, w)] -> [(a, w)]
forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' (f [(a, w)] -> f [(a, w)]) -> f [(a, w)] -> f [(a, w)]
forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b' f w
c' [(a, u)]
r' [(a, v)]
s'
    inner u
a v
_ f w
c [] ((a
y,v
b'):[(a, v)]
s') = let
      c' :: f w
c' = u -> v -> f w
p u
a v
b'
      in (w -> w -> [(a, w)] -> [(a, w)])
-> f w -> f w -> f [(a, w)] -> f [(a, w)]
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (a -> w -> w -> [(a, w)] -> [(a, w)]
forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
y) f w
c f w
c' (f [(a, w)] -> f [(a, w)]) -> f [(a, w)] -> f [(a, w)]
forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b' f w
c' [] [(a, v)]
s'
    inner u
_ v
b f w
c ((a
x,u
a'):[(a, u)]
r') [] = let
      c' :: f w
c' = u -> v -> f w
p u
a' v
b
      in (w -> w -> [(a, w)] -> [(a, w)])
-> f w -> f w -> f [(a, w)] -> f [(a, w)]
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (a -> w -> w -> [(a, w)] -> [(a, w)]
forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' (f [(a, w)] -> f [(a, w)]) -> f [(a, w)] -> f [(a, w)]
forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b f w
c' [(a, u)]
r' []
    inner u
_ v
_ f w
_ [] [] = [(a, w)] -> f [(a, w)]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []

    run :: Piecewise k u -> Piecewise k v -> f (Piecewise k w)
run (Piecewise u
a Map k u
f) (Piecewise v
b Map k v
g) = let
      c :: f w
c = u -> v -> f w
p u
a v
b
      l :: f [(k, w)]
l = u -> v -> f w -> [(k, u)] -> [(k, v)] -> f [(k, w)]
forall {a}.
Ord a =>
u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b f w
c (Map k u -> [(k, u)]
forall k a. Map k a -> [(k, a)]
M.toList Map k u
f) (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k v
g)
      in (w -> Map k w -> Piecewise k w)
-> f w -> f (Map k w) -> f (Piecewise k w)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 w -> Map k w -> Piecewise k w
forall k v. v -> Map k v -> Piecewise k v
Piecewise f w
c ([(k, w)] -> Map k w
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(k, w)] -> Map k w) -> f [(k, w)] -> f (Map k w)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f [(k, w)]
l)

    in Piecewise k u -> Piecewise k v -> f (Piecewise k w)
forall {k}.
Ord k =>
Piecewise k u -> Piecewise k v -> f (Piecewise k w)
run

instance Neighbourly (Piecewise k) where
  neighbours :: forall v. Ord v => Piecewise k v -> Set (v, v)
neighbours Piecewise k v
m = let
    pairs :: [b] -> [(b, b)]
pairs (b
x:r :: [b]
r@(b
y:[b]
_)) = (b
x,b
y)(b, b) -> [(b, b)] -> [(b, b)]
forall a. a -> [a] -> [a]
:[b] -> [(b, b)]
pairs [b]
r
    pairs [b]
_           = []
    in [(v, v)] -> Set (v, v)
forall a. Ord a => [a] -> Set a
S.fromList ([(v, v)] -> Set (v, v)) -> ([v] -> [(v, v)]) -> [v] -> Set (v, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> [(v, v)]
forall {b}. [b] -> [(b, b)]
pairs ([v] -> Set (v, v)) -> [v] -> Set (v, v)
forall a b. (a -> b) -> a -> b
$ Piecewise k v -> [v]
forall k v. Piecewise k v -> [v]
values Piecewise k v
m

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Semigroup b) => Semigroup (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Monoid b) => Monoid (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Num b) => Num (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Boolean b) => Boolean (Piecewise k b)

-- | Alter keys according to a function, assumed to be monotone (not checked)
mapKeysMonotonic :: (k -> l) -> Piecewise k v -> Piecewise l v
mapKeysMonotonic :: forall k l v. (k -> l) -> Piecewise k v -> Piecewise l v
mapKeysMonotonic k -> l
f (Piecewise v
a Map k v
m) = v -> Map l v -> Piecewise l v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a ((k -> l) -> Map k v -> Map l v
forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeysMonotonic k -> l
f Map k v
m)

-- | Alter keys according to a function, assumed to be antitone (not checked)
mapKeysAntitonic :: (k -> l) -> Piecewise k v -> Piecewise l v
mapKeysAntitonic :: forall k l v. (k -> l) -> Piecewise k v -> Piecewise l v
mapKeysAntitonic k -> l
f = let

  inner :: a -> Map k a -> (a, Map l a)
inner a
a Map k a
MI.Tip = (a
a, Map l a
forall k a. Map k a
MI.Tip)
  inner a
a (MI.Bin Int
s k
x a
b Map k a
l Map k a
r) = let
    (a
a', Map l a
l') = a -> Map k a -> (a, Map l a)
inner a
a Map k a
l
    (a
b', Map l a
r') = a -> Map k a -> (a, Map l a)
inner a
b Map k a
r
    in (a
b', Int -> l -> a -> Map l a -> Map l a -> Map l a
forall k a. Int -> k -> a -> Map k a -> Map k a -> Map k a
MI.Bin Int
s (k -> l
f k
x) a
a' Map l a
r' Map l a
l')

  start :: Piecewise k a -> Piecewise l a
start (Piecewise a
a Map k a
m) = (a -> Map l a -> Piecewise l a) -> (a, Map l a) -> Piecewise l a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Map l a -> Piecewise l a
forall k v. v -> Map k v -> Piecewise k v
Piecewise ((a, Map l a) -> Piecewise l a) -> (a, Map l a) -> Piecewise l a
forall a b. (a -> b) -> a -> b
$ a -> Map k a -> (a, Map l a)
forall {a}. a -> Map k a -> (a, Map l a)
inner a
a Map k a
m
  in Piecewise k v -> Piecewise l v
forall {a}. Piecewise k a -> Piecewise l a
start

-- | Split in two: one which assumes keys are less than `k`, and one
-- which assumes them greater than or equal to `k`.
splitPiecewise :: Ord k => k -> Piecewise k v -> (Piecewise k v, Piecewise k v)
splitPiecewise :: forall k v.
Ord k =>
k -> Piecewise k v -> (Piecewise k v, Piecewise k v)
splitPiecewise k
k (Piecewise v
a Map k v
m) = case k -> Map k v -> (Map k v, Maybe v, Map k v)
forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
M.splitLookup k
k Map k v
m of
  (Map k v
m1, Just v
b, Map k v
m2) -> (v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a Map k v
m1, v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
b Map k v
m2)
  (Map k v
m1, Maybe v
Nothing, Map k v
m2) -> let
    p1 :: Piecewise k v
p1 = v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a Map k v
m1
    in (Piecewise k v
p1, v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise (Piecewise k v -> v
forall k v. Piecewise k v -> v
rightEnd Piecewise k v
p1) Map k v
m2)

-- | Assemble two maps; it is assumed that all keys in the left-hand
-- map are less than `k` and all keys in the right-hand map are
-- greater than or equal to `k` (which is not checked).
gluePiecewise :: (Eq v) => Piecewise k v -> k -> Piecewise k v -> Piecewise k v
gluePiecewise :: forall v k.
Eq v =>
Piecewise k v -> k -> Piecewise k v -> Piecewise k v
gluePiecewise p :: Piecewise k v
p@(Piecewise v
a Map k v
m) k
k (Piecewise v
c Map k v
n) = let
  b :: v
b = Piecewise k v -> v
forall k v. Piecewise k v -> v
rightEnd Piecewise k v
p
  in v -> Map k v -> Piecewise k v
forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a (if v
b v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
c then Map k v -> Map k v -> Map k v
forall k a. Map k a -> Map k a -> Map k a
MI.link2 Map k v
m Map k v
n else k -> v -> Map k v -> Map k v -> Map k v
forall k a. k -> a -> Map k a -> Map k a -> Map k a
MI.link k
k v
c Map k v
m Map k v
n)

-- | This is almost a monad (with `cst` as `pure`) except that we need
-- an `Eq` instance on the values.
mjoin :: (Ord k, Eq w) => (v -> Piecewise k w) -> Piecewise k v -> Piecewise k w
mjoin :: forall k w v.
(Ord k, Eq w) =>
(v -> Piecewise k w) -> Piecewise k v -> Piecewise k w
mjoin v -> Piecewise k w
f (Piecewise v
a Map k v
m) = let
  inner :: Piecewise a v -> [(a, Piecewise a v)] -> Piecewise a v
inner Piecewise a v
p []        = Piecewise a v
p
  inner Piecewise a v
p ((a
k,Piecewise a v
q):[(a, Piecewise a v)]
l) = let
    (Piecewise a v
p',  Piecewise a v
_) = a -> Piecewise a v -> (Piecewise a v, Piecewise a v)
forall k v.
Ord k =>
k -> Piecewise k v -> (Piecewise k v, Piecewise k v)
splitPiecewise a
k Piecewise a v
p
    (Piecewise a v
_ , Piecewise a v
q') = a -> Piecewise a v -> (Piecewise a v, Piecewise a v)
forall k v.
Ord k =>
k -> Piecewise k v -> (Piecewise k v, Piecewise k v)
splitPiecewise a
k Piecewise a v
q
    in Piecewise a v -> a -> Piecewise a v -> Piecewise a v
forall v k.
Eq v =>
Piecewise k v -> k -> Piecewise k v -> Piecewise k v
gluePiecewise Piecewise a v
p' a
k (Piecewise a v -> Piecewise a v) -> Piecewise a v -> Piecewise a v
forall a b. (a -> b) -> a -> b
$ Piecewise a v -> [(a, Piecewise a v)] -> Piecewise a v
inner Piecewise a v
q' [(a, Piecewise a v)]
l
  in Piecewise k w -> [(k, Piecewise k w)] -> Piecewise k w
forall {a} {v}.
(Ord a, Eq v) =>
Piecewise a v -> [(a, Piecewise a v)] -> Piecewise a v
inner (v -> Piecewise k w
f v
a) ((v -> Piecewise k w) -> (k, v) -> (k, Piecewise k w)
forall a b. (a -> b) -> (k, a) -> (k, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap v -> Piecewise k w
f ((k, v) -> (k, Piecewise k w)) -> [(k, v)] -> [(k, Piecewise k w)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k v
m)