-- SPDX-FileCopyrightText: Copyright Preetham Gujjula
-- SPDX-License-Identifier: BSD-3-Clause
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module: Data.List.ApplyMerge
-- License: BSD-3-Clause
-- Maintainer: Preetham Gujjula <libraries@mail.preetham.io>
-- Stability: experimental
module Data.List.ApplyMerge (applyMerge, applyMergeBy, applyMergeOn) where

import ApplyMerge.IntSet qualified
import Data.Proxy (Proxy (..))
import Data.Reflection (Reifies, reflect, reify)
import Data.Semigroup (Arg (..))

-- | If given a binary function @f@ that is non-decreasing in both arguments,
--   and two (potentially infinite) ordered lists @xs@ and @ys@, then
--   @'applyMerge' f xs ys@ is a sorted list of all @f x y@, for each @x@ in
--   @xs@ and @y@ in @ys@.
--
--   Producing \(n\) elements of @'applyMerge' f xs ys@ takes \(O(n \log n)\)
--   time and \(O(\sqrt{n})\) auxiliary space, assuming that @f@ and @compare@
--   take \(O(1)\) time.
--
--   For example, to generate the 3-smooth numbers
--   ([Wikipedia](https://en.wikipedia.org/wiki/Smooth_number)):
--
--   > smooth3 :: [Integer]
--   > smooth3 = applyMerge (*) (iterate (*2) 1) (iterate (*3) 1)
--
--   For more examples, see
--   [README#examples](https://github.com/pgujjula/apply-merge/#examples).
applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge :: forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge = (a -> b -> c) -> [a] -> [b] -> [c]
forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
ApplyMerge.IntSet.applyMerge

-- | Like 'applyMerge', but uses a custom comparison function.
applyMergeBy :: (c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeBy :: forall c a b.
(c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeBy = (c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c.
(c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeBy_

-- Reflection logic in applyMerge_ is based on "All about reflection: a
-- tutorial" [1] by Arnaud Spiwack, licensed under CC BY 4.0 [2].
--
-- [1]: https://www.tweag.io/blog/2017-12-21-reflection-tutorial/
-- [2]: https://creativecommons.org/licenses/by/4.0/
applyMergeBy_ ::
  forall a b c. (c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeBy_ :: forall a b c.
(c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeBy_ c -> c -> Ordering
cmp a -> b -> c
f [a]
as [b]
bs =
  (c -> c -> Ordering)
-> (forall {s}. Reifies s (c -> c -> Ordering) => Proxy s -> [c])
-> [c]
forall a r. a -> (forall s. Reifies s a => Proxy s -> r) -> r
reify c -> c -> Ordering
cmp ((forall {s}. Reifies s (c -> c -> Ordering) => Proxy s -> [c])
 -> [c])
-> (forall {s}. Reifies s (c -> c -> Ordering) => Proxy s -> [c])
-> [c]
forall a b. (a -> b) -> a -> b
$ \(Proxy s
_ :: Proxy s) ->
    let f' :: a -> b -> ReflectedOrd s c
        f' :: a -> b -> ReflectedOrd s c
f' a
a b
b = c -> ReflectedOrd s c
forall {k} (s :: k) a. a -> ReflectedOrd s a
ReflectedOrd (a -> b -> c
f a
a b
b)
     in (ReflectedOrd s c -> c) -> [ReflectedOrd s c] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map ReflectedOrd s c -> c
forall {k} (s :: k) a. ReflectedOrd s a -> a
unReflectedOrd ((a -> b -> ReflectedOrd s c) -> [a] -> [b] -> [ReflectedOrd s c]
forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge a -> b -> ReflectedOrd s c
f' [a]
as [b]
bs)

newtype ReflectedOrd s a = ReflectedOrd {forall {k} (s :: k) a. ReflectedOrd s a -> a
unReflectedOrd :: a}

instance (Reifies s (a -> a -> Ordering)) => Eq (ReflectedOrd s a) where
  == :: ReflectedOrd s a -> ReflectedOrd s a -> Bool
(==) (ReflectedOrd a
x) (ReflectedOrd a
y) =
    let cmp :: a -> a -> Ordering
cmp = Proxy s -> a -> a -> Ordering
forall {k} (s :: k) a (proxy :: k -> *).
Reifies s a =>
proxy s -> a
forall (proxy :: k -> *). proxy s -> a -> a -> Ordering
reflect (Proxy s
forall {k} (t :: k). Proxy t
Proxy :: Proxy s)
     in a -> a -> Ordering
cmp a
x a
y Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ

instance (Reifies s (a -> a -> Ordering)) => Ord (ReflectedOrd s a) where
  compare :: ReflectedOrd s a -> ReflectedOrd s a -> Ordering
compare (ReflectedOrd a
x) (ReflectedOrd a
y) =
    let cmp :: a -> a -> Ordering
cmp = Proxy s -> a -> a -> Ordering
forall {k} (s :: k) a (proxy :: k -> *).
Reifies s a =>
proxy s -> a
forall (proxy :: k -> *). proxy s -> a -> a -> Ordering
reflect (Proxy s
forall {k} (t :: k). Proxy t
Proxy :: Proxy s)
     in a -> a -> Ordering
cmp a
x a
y

-- | Like 'applyMerge', but applies a custom projection function before
--   performing comparisons.
--
--   For example, to compute the Gaussian integers, ordered by norm:
--
--   > zs :: [Integer]
--   > zs = 0 : concatMap (\i -> [i, -i]) [1..]
--   >
--   > gaussianIntegers :: [GaussianInteger]      -- `GaussianInteger` from arithmoi
--   > gaussianIntegers = applyMergeOn norm (:+) zs zs
applyMergeOn ::
  (Ord d) => (c -> d) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeOn :: forall d c a b.
Ord d =>
(c -> d) -> (a -> b -> c) -> [a] -> [b] -> [c]
applyMergeOn c -> d
p a -> b -> c
f [a]
as [b]
bs =
  let f' :: a -> b -> Arg d c
f' a
a b
b =
        let c :: c
c = a -> b -> c
f a
a b
b
         in d -> c -> Arg d c
forall a b. a -> b -> Arg a b
Arg (c -> d
p c
c) c
c
   in (Arg d c -> c) -> [Arg d c] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map (\(Arg d
_ c
c) -> c
c) ((a -> b -> Arg d c) -> [a] -> [b] -> [Arg d c]
forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge a -> b -> Arg d c
f' [a]
as [b]
bs)