{-# LANGUAGE Safe #-}

module Invert
  ( -- * Overview
    -- $overview

    -- * 1. Varieties of function
    function,
    bijection,
    injection,
    surjection,

    -- * 2. Inversion strategies
    linearSearchLazy,
    linearSearchStrict,
    binarySearch,
    hashTable,

    -- * 3. Domain enumeration
    enumBounded,
    genum,

    -- * The Strategy type
    Strategy,
    -- $strategyCreation
    strategyAll,
    strategyOneAndAll,

    -- * Re-exports
    -- $reexports
    module Invert.Reexport,
  )
where

import Data.Eq (Eq, (==))
import Data.Foldable (foldl')
import Data.Function ((.))
import Data.List qualified as List (lookup, map)
import Data.List.NonEmpty (NonEmpty, nonEmpty)
import Data.Maybe (Maybe (Just, Nothing), fromMaybe, listToMaybe)
import Data.Maybe qualified as List (mapMaybe)
import Data.Ord (Ord)
import Data.Tuple (uncurry)
import Generics.Deriving qualified as GEnum (genum)
import Invert.Reexport
import Map (Map (Map))
import Map qualified
import Vector qualified
import Prelude (Bounded, Enum, enumFromTo, error, maxBound, minBound)

-- $overview
--
-- There are three considerations when you’re inverting a function:
--
--   1. Is it an injection, a surjection, both (a bijection), or neither?
--   2. What data structure do you want to use for efficient lookups?
--   3. Can you produce a list of all values in the function’s domain?
--
-- === 1. What sort of function do you have?
--
-- This question determines the type of the function’s inverse.
--
-- For a function @(a -> b)@, we call @(a)@ its /domain/, and @(b)@ its /codomain/.
--
--   * In general, when you invert a 'function' of type @(a -> b)@,
--     the type of the inverse is @(b -> [a])@.
--     The result is a list because it contains all domain values that
--     map to a given codomain value; there may be none, one, or many.
--
--   * If your function @(a -> b)@ is a 'bijection',
--     you can invert it to get a function @(b -> a)@.
--     Bijections are quite pleasing in this way.
--
--   * If no two domain values map to the same codomain value,
--     then your function is an 'injection',
--     and it has an inverse of type @(b -> 'Maybe' a)@.
--
--   * If every codomain value has some domain value that maps to it,
--     then your function is a 'surjection',
--     and it has an inverse of type @(b -> 'NonEmpty' a)@.
--
-- You are responsible for determining which is appropriate for a particular
-- situation: 'function', 'bijection', 'injection', or 'surjection'.
-- Choose carefully; the wrong choice may produce an inverse which is
-- partial or incorrect.
--
-- === 2. How can we produce a reasonably efficient inversion?
--
-- The simplest inversion strategies, 'linearSearchLazy' and 'linearSearchStrict',
-- apply the function to each element of the domain, one by one.
-- We call this a /linear search/ because the time required for each
-- application has a linear correspondence with the size of the domain.
--
--   * 'linearSearchStrict' works by precomputing a strict sequence
--     of tuples, one for each value of the domain.
--
--   * 'linearSearchLazy' precomputes nothing at all.
--     It is possible to use this strategy when the domain is infinite.
--
-- Our other two strategies, 'binarySearch' and 'hashTable',
-- work by building data structures that allow more efficient lookups.
--
--   * 'binarySearch' precomputes a binary search tree;
--     the codomain must belong to the 'Ord' class.
--
--   * 'hashTable' precomputes a hash table;
--     the codomain must belong to the 'Hashable' class.
--
-- The 'Hashable' class comes from "Data.Hashable" in the @hashable@ package.
-- The class is re-exported by "Invert", which you may find convenient if
-- your primary motivation for deriving 'Hashable' is to invert a function.
--
-- === 3. How will you enumerate the domain?
--
-- Inverting a function @(a -> b)@ requires having a list of all
-- possible values of domain @(a)@; from this, we can apply the
-- function to every value to produce a list of tuples that
-- completely describes the function.
--
-- We offer two suggestions for automatically producing this list:
--
--   * 'enumBounded' uses two stock-derivable classes, 'Enum' and 'Bounded'.
--   * 'genum' uses GHC generics; it requires deriving 'Generic' and 'GEnum'.
--
-- The 'Generic' class comes from "GHC.Generics", and the 'GEnum' class
-- comes from "Generics.Deriving" in the @generic-deriving@ package.
-- Both classes are re-exported by "Invert", which you may find convenient
-- if your primary motivation for deriving 'GEnum' is to invert a function.

function ::
  Strategy a b ->
  -- | A complete list of all the values of the domain.
  [a] ->
  -- | The function to invert.
  (a -> b) ->
  -- | The inverse of the given function.
  (b -> [a])
bijection ::
  Strategy a b ->
  -- | A complete list of all the values of the domain.
  [a] ->
  -- | The function to invert.
  --   __This function must be bijective!__
  --   This means that every value in the codomain has
  --   exactly one value in the domain that maps to it.
  (a -> b) ->
  -- | The inverse of the given function.
  (b -> a)
injection ::
  Strategy a b ->
  -- | A complete list of all the values of the domain.
  [a] ->
  -- | The function to invert.
  --   __This function must be injective!__
  --   This means that no two values in the domain map
  --   to the same value of the codomain.
  (a -> b) ->
  -- | The inverse of the given function.
  (b -> Maybe a)
surjection ::
  Strategy a b ->
  -- | A complete list of all the values of the domain.
  [a] ->
  -- | The function to invert.
  --   __This function must be surjective!__
  --   This means that every value in the codomain has
  --   at least one value in the domain that maps to it.
  (a -> b) ->
  -- | The inverse of the given function.
  (b -> NonEmpty a)
function :: forall a b. Strategy a b -> [a] -> (a -> b) -> b -> [a]
function (Strategy [(b, a)] -> b -> Maybe a
_ [(b, a)] -> b -> [a]
s) [a]
as a -> b
f = [(b, a)] -> b -> [a]
s ([a] -> (a -> b) -> [(b, a)]
forall a b. [a] -> (a -> b) -> [(b, a)]
inverseEntries [a]
as a -> b
f)

injection :: forall a b. Strategy a b -> [a] -> (a -> b) -> b -> Maybe a
injection (Strategy [(b, a)] -> b -> Maybe a
s [(b, a)] -> b -> [a]
_) [a]
as a -> b
f = [(b, a)] -> b -> Maybe a
s ([a] -> (a -> b) -> [(b, a)]
forall a b. [a] -> (a -> b) -> [(b, a)]
inverseEntries [a]
as a -> b
f)

bijection :: forall a b. Strategy a b -> [a] -> (a -> b) -> b -> a
bijection (Strategy [(b, a)] -> b -> Maybe a
s [(b, a)] -> b -> [a]
_) [a]
as a -> b
f = Maybe a -> a
forall {a}. Maybe a -> a
finagle (Maybe a -> a) -> (b -> Maybe a) -> b -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(b, a)] -> b -> Maybe a
s ([a] -> (a -> b) -> [(b, a)]
forall a b. [a] -> (a -> b) -> [(b, a)]
inverseEntries [a]
as a -> b
f)
  where
    finagle :: Maybe a -> a
finagle = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Not a bijection!")

surjection :: forall a b. Strategy a b -> [a] -> (a -> b) -> b -> NonEmpty a
surjection (Strategy [(b, a)] -> b -> Maybe a
_ [(b, a)] -> b -> [a]
s) [a]
as a -> b
f = [a] -> NonEmpty a
forall {a}. [a] -> NonEmpty a
finagle ([a] -> NonEmpty a) -> (b -> [a]) -> b -> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(b, a)] -> b -> [a]
s ([a] -> (a -> b) -> [(b, a)]
forall a b. [a] -> (a -> b) -> [(b, a)]
inverseEntries [a]
as a -> b
f)
  where
    finagle :: [a] -> NonEmpty a
finagle = NonEmpty a -> Maybe (NonEmpty a) -> NonEmpty a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> NonEmpty a
forall a. HasCallStack => [Char] -> a
error [Char]
"Not a surjection!") (Maybe (NonEmpty a) -> NonEmpty a)
-> ([a] -> Maybe (NonEmpty a)) -> [a] -> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty

-- | An inversion strategy is an approach for producing
--    the inverse of an @(a -> b)@ function
--
-- All strategies produce the same results, but they
-- have operational differences that affect performance.
data Strategy a b
  = Strategy
      ([(b, a)] -> b -> Maybe a)
      ([(b, a)] -> b -> [a])

-- $strategyCreation
--
-- === Defining your own strategies
--
-- If you want to design your own strategy instead
-- of using one provided by this module, use either
-- 'strategyAll' or 'strategyOneAndAll'.

strategyAll ::
  -- | Find all matches
  ([(b, a)] -> b -> [a]) ->
  Strategy a b
strategyAll :: forall b a. ([(b, a)] -> b -> [a]) -> Strategy a b
strategyAll [(b, a)] -> b -> [a]
all = ([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
forall b a.
([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
strategyOneAndAll [(b, a)] -> b -> Maybe a
one [(b, a)] -> b -> [a]
all
  where
    one :: [(b, a)] -> b -> Maybe a
one [(b, a)]
bas b
b = [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe ([(b, a)] -> b -> [a]
all [(b, a)]
bas b
b)

strategyOneAndAll ::
  -- | Find the first match
  ([(b, a)] -> b -> Maybe a) ->
  -- | Find all matches
  ([(b, a)] -> b -> [a]) ->
  Strategy a b
strategyOneAndAll :: forall b a.
([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
strategyOneAndAll = ([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
forall a b.
([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
Strategy

inverseEntries :: [a] -> (a -> b) -> [(b, a)]
inverseEntries :: forall a b. [a] -> (a -> b) -> [(b, a)]
inverseEntries [a]
as a -> b
f = (a -> (b, a)) -> [a] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
List.map (\a
a -> (a -> b
f a
a, a
a)) [a]
as

mapStrategy :: Map Maybe b a -> Map [] b a -> Strategy a b
mapStrategy :: forall b a. Map Maybe b a -> Map [] b a -> Strategy a b
mapStrategy Map Maybe b a
one Map [] b a
all = ([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
forall a b.
([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
Strategy (Map Maybe b a -> [(b, a)] -> b -> Maybe a
forall {f :: * -> *} {a} {b}. Map f a b -> [(a, b)] -> a -> f b
f Map Maybe b a
one) (Map [] b a -> [(b, a)] -> b -> [a]
forall {f :: * -> *} {a} {b}. Map f a b -> [(a, b)] -> a -> f b
f Map [] b a
all)
  where
    f :: Map f a b -> [(a, b)] -> a -> f b
f Map {map
empty :: map
empty :: ()
Map.empty, a -> b -> map
singleton :: a -> b -> map
singleton :: ()
Map.singleton, map -> map -> map
union :: map -> map -> map
union :: ()
Map.union, map -> a -> f b
lookup :: map -> a -> f b
lookup :: ()
Map.lookup} =
      map -> a -> f b
lookup (map -> a -> f b) -> ([(a, b)] -> map) -> [(a, b)] -> a -> f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (map -> map -> map) -> map -> [map] -> map
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' map -> map -> map
union map
empty ([map] -> map) -> ([(a, b)] -> [map]) -> [(a, b)] -> map
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> map) -> [(a, b)] -> [map]
forall a b. (a -> b) -> [a] -> [b]
List.map ((a -> b -> map) -> (a, b) -> map
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> map
singleton)

-- | A function inversion strategy that precomputes nothing at all
--
-- It is possible to use this strategy when the domain is infinite.
linearSearchLazy :: Eq b => Strategy a b
linearSearchLazy :: forall b a. Eq b => Strategy a b
linearSearchLazy = ([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
forall a b.
([(b, a)] -> b -> Maybe a)
-> ([(b, a)] -> b -> [a]) -> Strategy a b
Strategy [(b, a)] -> b -> Maybe a
forall {a} {b}. Eq a => [(a, b)] -> a -> Maybe b
one [(b, a)] -> b -> [a]
forall {b} {b}. Eq b => [(b, b)] -> b -> [b]
all
  where
    one :: [(a, b)] -> a -> Maybe b
one [(a, b)]
bas a
b = a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
List.lookup a
b [(a, b)]
bas
    all :: [(b, b)] -> b -> [b]
all [(b, b)]
bas b
b = ((b, b) -> Maybe b) -> [(b, b)] -> [b]
forall a b. (a -> Maybe b) -> [a] -> [b]
List.mapMaybe (b -> (b, b) -> Maybe b
forall b a. Eq b => b -> (b, a) -> Maybe a
sndIfFstEq b
b) [(b, b)]
bas

-- | A function inversion strategy that works by precomputing a
--    strict sequence of tuples, one for each value of the domain
--
-- For larger functions, it may be preferable to use 'binarySearch' or
-- 'hashTable' instead to get a more efficient inverse.
linearSearchStrict :: Eq b => Strategy a b
linearSearchStrict :: forall b a. Eq b => Strategy a b
linearSearchStrict = ([(b, a)] -> b -> [a]) -> Strategy a b
forall b a. ([(b, a)] -> b -> [a]) -> Strategy a b
strategyAll [(b, a)] -> b -> [a]
forall {b} {b}. Eq b => [(b, b)] -> b -> [b]
f
  where
    f :: [(b, a)] -> b -> [a]
f [(b, a)]
bas b
b = Vector a -> [a]
forall a. Vector a -> [a]
Vector.toList (((b, a) -> Maybe a) -> Vector (b, a) -> Vector a
forall a b. (a -> Maybe b) -> Vector a -> Vector b
Vector.mapMaybe (b -> (b, a) -> Maybe a
forall b a. Eq b => b -> (b, a) -> Maybe a
sndIfFstEq b
b) Vector (b, a)
v)
      where
        v :: Vector (b, a)
v = [(b, a)] -> Vector (b, a)
forall a. [a] -> Vector a
Vector.fromList [(b, a)]
bas

sndIfFstEq :: Eq b => b -> (b, a) -> Maybe a
sndIfFstEq :: forall b a. Eq b => b -> (b, a) -> Maybe a
sndIfFstEq b
x (b
b, a
a) = if b
b b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
x then a -> Maybe a
forall a. a -> Maybe a
Just a
a else Maybe a
forall a. Maybe a
Nothing

-- | A function inversion strategy that works by precomputing
--    a binary search tree
--
-- The data structure imposes the requirement that the codomain
-- belongs to the 'Ord' class.
binarySearch :: Ord b => Strategy a b
binarySearch :: forall b a. Ord b => Strategy a b
binarySearch = Map Maybe b a -> Map [] b a -> Strategy a b
forall b a. Map Maybe b a -> Map [] b a -> Strategy a b
mapStrategy Map Maybe b a
forall a b. Ord a => SingleMap a b
Map.ordSingleMap Map [] b a
forall a b. Ord a => MultiMap a b
Map.ordMultiMap

-- | A function inversion strategy that works by precomputing
--    a hash table
--
-- The data structure imposes the requirement that the codomain
-- belongs to the 'Hashable' class.
hashTable :: (Eq b, Hashable b) => Strategy a b
hashTable :: forall b a. (Eq b, Hashable b) => Strategy a b
hashTable = Map Maybe b a -> Map [] b a -> Strategy a b
forall b a. Map Maybe b a -> Map [] b a -> Strategy a b
mapStrategy Map Maybe b a
forall a b. (Eq a, Hashable a) => SingleMap a b
Map.hashSingleMap Map [] b a
forall a b. (Eq a, Hashable a) => MultiMap a b
Map.hashMultiMap

-- | A convenient way to enumerate the domain for a function that you
-- want to invert, using the stock-derivable classes 'Enum' and 'Bounded'
--
-- To derive the required typeclass instances, add the following deriving clause to
-- the type’s definition:
--
-- @
-- deriving (Enum, Bounded)
-- @
enumBounded :: (Enum a, Bounded a) => [a]
enumBounded :: forall a. (Enum a, Bounded a) => [a]
enumBounded = a -> a -> [a]
forall a. Enum a => a -> a -> [a]
enumFromTo a
forall a. Bounded a => a
minBound a
forall a. Bounded a => a
maxBound

-- | Use GHC generics to enumerate a function's domain
--
-- This requires deriving 'Generic' and 'GEnum'. The 'Generic' class comes
-- from "GHC.Generics", and the 'GEnum' class comes from "Generics.Deriving"
-- in the @generic-deriving@ package.
--
-- To derive the required typeclass instances, enable the following
-- language extensions:
--
-- @
-- \{\-# language DeriveGeneric, DeriveAnyClass, DerivingStrategies #\-\}
-- @
--
-- Then add the following deriving clauses to the type’s definition:
--
-- @
-- deriving stock Generic
-- deriving anyclass GEnum
-- @
genum :: GEnum a => [a]
genum :: forall a. GEnum a => [a]
genum = [a]
forall a. GEnum a => [a]
GEnum.genum

-- $reexports
--
-- This module provides a few definitions that come directly from
-- other packages. These are here to let you conveniently derive
-- 'Hashable' and 'GEnum' with only the "Invert" module imported.
--
-- List of re-exports:
--
--   - __'Hashable'__ (for the 'hashTable' inversion strategy)
--   - __'Generic'__ and __'GEnum'__ (for the 'genum' domain
--     enumeration approach)