{-# LANGUAGE Safe #-}
module Data.Universe.Helpers (
  -- | This module is for functions that are useful for writing instances,
  -- but not necessarily for using them (and hence are not exported by the
  -- main module to avoid cluttering up the namespace).

  -- * Building lists
  universeDef,
  interleave,
  diagonal,
  diagonals,
  (+++),
  cartesianProduct,
  (+*+),
  (<+*+>),
  choices,

  -- * Building cardinalities
  -- | These functions are handy for inheriting the definition of
  -- 'Data.Universe.Class.cardinality' in a newtype instance. For example,
  -- one might write
  --
  -- > newtype Foo = Foo Bar
  -- > instance Finite Foo where cardinality = retagWith Foo cardinality
  retagWith,
  retag,
  Tagged (..),
  Natural,

  -- * Debugging
  -- | These functions exist primarily as a specification to test against.
  unfairCartesianProduct,
  unfairChoices
  ) where

import Data.List
import Data.Tagged (Tagged (..), retag)
import Numeric.Natural (Natural)

-- | For many types, the 'universe' should be @[minBound .. maxBound]@;
-- 'universeDef' makes it easy to make such types an instance of 'Universe' via
-- the snippet
--
-- > instance Universe Foo where universe = universeDef
universeDef :: (Bounded a, Enum a) => [a]
universeDef :: forall a. (Bounded a, Enum a) => [a]
universeDef = [a
forall a. Bounded a => a
minBound .. a
forall a. Bounded a => a
maxBound]

-- | Fair n-way interleaving: given a finite number of (possibly infinite)
-- lists, produce a single list such that whenever @v@ has finite index in one
-- of the input lists, @v@ also has finite index in the output list. No list's
-- elements occur more frequently (on average) than another's.
interleave :: [[a]] -> [a]
interleave :: forall a. [[a]] -> [a]
interleave = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([[a]] -> [[a]]) -> [[a]] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose

-- | Unfair n-way interleaving: given a possibly infinite number of (possibly
-- infinite) lists, produce a single list such that whenever @v@ has finite
-- index in an input list at finite index, @v@ also has finite index in the
-- output list. Elements from lists at lower index occur more frequently, but
-- not exponentially so.
diagonal :: [[a]] -> [a]
diagonal :: forall a. [[a]] -> [a]
diagonal = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([[a]] -> [[a]]) -> [[a]] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
diagonals

-- | Like 'diagonal', but expose a tiny bit more (non-semantic) information:
-- if you lay out the input list in two dimensions, each list in the result
-- will be one of the diagonals of the input. In particular, each element of
-- the output will be a list whose elements are each from a distinct input
-- list.
diagonals :: [[a]] -> [[a]]
diagonals :: forall a. [[a]] -> [[a]]
diagonals = [[a]] -> [[a]]
forall a. HasCallStack => [a] -> [a]
tail ([[a]] -> [[a]]) -> ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]] -> [[a]]
forall {a}. [[a]] -> [[a]] -> [[a]]
go [] where
  -- it is critical for some applications that we start producing answers
  -- before inspecting es_
  go :: [[a]] -> [[a]] -> [[a]]
go [[a]]
b [[a]]
es_ = [a
h | a
h:[a]
_ <- [[a]]
b] [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: case [[a]]
es_ of
    []   -> [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose [[a]]
ts
    [a]
e:[[a]]
es -> [[a]] -> [[a]] -> [[a]]
go ([a]
e[a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
:[[a]]
ts) [[a]]
es
    where ts :: [[a]]
ts = [[a]
t | a
_:[a]
t <- [[a]]
b]

-- | Fair 2-way interleaving.
(+++) :: [a] -> [a] -> [a]
[a]
xs +++ :: forall a. [a] -> [a] -> [a]
+++ [a]
ys = [[a]] -> [a]
forall a. [[a]] -> [a]
interleave [[a]
xs,[a]
ys]

-- | Slightly unfair 2-way Cartesian product: given two (possibly infinite)
-- lists, produce a single list such that whenever @v@ and @w@ have finite
-- indices in the input lists, @(v,w)@ has finite index in the output list.
-- Lower indices occur as the @fst@ part of the tuple more frequently, but not
-- exponentially so.
cartesianProduct :: (a -> b -> c) -> [a] -> [b] -> [c]
-- special case: don't want to construct an infinite list of empty lists to pass to diagonal
cartesianProduct :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
cartesianProduct a -> b -> c
_ []   [b]
_  = []
cartesianProduct a -> b -> c
f [a]
xs  [b]
ys  = [[c]] -> [c]
forall a. [[a]] -> [a]
diagonal [[a -> b -> c
f a
x b
y | a
x <- [a]
xs] | b
y <- [b]
ys]

-- | @'cartesianProduct' (,)@
(+*+) :: [a] -> [b] -> [(a,b)]
+*+ :: forall a b. [a] -> [b] -> [(a, b)]
(+*+) = (a -> b -> (a, b)) -> [a] -> [b] -> [(a, b)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
cartesianProduct (,)

-- | A '+*+' with application.
--
-- @'cartesianProduct' ($)@
(<+*+>) :: [a -> b] -> [a] -> [b]
<+*+> :: forall a b. [a -> b] -> [a] -> [b]
(<+*+>) = ((a -> b) -> a -> b) -> [a -> b] -> [a] -> [b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
cartesianProduct (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)

-- | Slightly unfair n-way Cartesian product: given a finite number of
-- (possibly infinite) lists, produce a single list such that whenever @vi@ has
-- finite index in list i for each i, @[v1, ..., vn]@ has finite index in the
-- output list.
choices :: [[a]] -> [[a]]
choices :: forall a. [[a]] -> [[a]]
choices = ([a] -> [[a]] -> [[a]]) -> [[a]] -> [[a]] -> [[a]]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((a -> [a] -> [a]) -> [a] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
cartesianProduct (:)) [[]]

retagWith :: (a -> b) -> Tagged a x -> Tagged b x
retagWith :: forall a b x. (a -> b) -> Tagged a x -> Tagged b x
retagWith a -> b
_ (Tagged x
n) = x -> Tagged b x
forall {k} (s :: k) b. b -> Tagged s b
Tagged x
n

-- | Very unfair 2-way Cartesian product: same guarantee as the slightly unfair
-- one, except that lower indices may occur as the @fst@ part of the tuple
-- exponentially more frequently.
unfairCartesianProduct :: (a -> b -> c) -> [a] -> [b] -> [c]
unfairCartesianProduct :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
unfairCartesianProduct a -> b -> c
_ [a]
_  [] = [] -- special case: don't want to walk down xs forever hoping one of them will produce a nonempty thing
unfairCartesianProduct a -> b -> c
f [a]
xs [b]
ys = [a] -> [b] -> [c]
go [a]
xs [b]
ys where
  go :: [a] -> [b] -> [c]
go (a
x:[a]
xs) [b]
ys = (b -> c) -> [b] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map (a -> b -> c
f a
x) [b]
ys [c] -> [c] -> [c]
forall a. [a] -> [a] -> [a]
+++ [a] -> [b] -> [c]
go [a]
xs [b]
ys
  go []     [b]
ys = []

-- | Very unfair n-way Cartesian product: same guarantee as the slightly unfair
-- one, but not as good in the same sense that the very unfair 2-way product is
-- worse than the slightly unfair 2-way product.
unfairChoices :: [[a]] -> [[a]]
unfairChoices :: forall a. [[a]] -> [[a]]
unfairChoices = ([a] -> [[a]] -> [[a]]) -> [[a]] -> [[a]] -> [[a]]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((((a, [a]) -> [a]) -> [(a, [a])] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> [a] -> [a]) -> (a, [a]) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (:)) ([(a, [a])] -> [[a]]) -> ([[a]] -> [(a, [a])]) -> [[a]] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) (([[a]] -> [(a, [a])]) -> [[a]] -> [[a]])
-> ([a] -> [[a]] -> [(a, [a])]) -> [a] -> [[a]] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [a] -> (a, [a])) -> [a] -> [[a]] -> [(a, [a])]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
unfairCartesianProduct (,)) [[]]