{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE RoleAnnotations     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE TypeOperators #-}
module Data.NatMap (
    NatMap(),

    -- * Entry
    Entry(),
    getKeyValue,
    makeEntry,
    makeIdEntry,
    unsafeMakeEntry,

    -- * Construction
    empty, singleton,
    partialIdentity,
    fromEntries,
    insert, delete,
    
    -- * Queries
    size,
    member, notMember,
    lookup, lookup_,
    keys,
    toEntries,

    -- * Map, Filter, Traversal

    map1, mapMaybe1, traverse1, wither1,
    mapWithKey1, mapMaybeWithKey1, traverseWithKey1, witherWithKey1,

    -- * Combinators

    union, unionWith, consistentUnion,

    -- * As partial natural transformation

    identity, compose,
    outerNat, innerNat, horizontalCompose,

    -- * Totality

    fullSize, isTotal,
    toTotal,

    -- * Re-exports
    (:~>)(..), wrapNT, unwrapNT,

    -- * Utility
    Var(), indices,
) where

import Prelude hiding (lookup)

import Data.Kind ( Type )

import Data.Proxy ( Proxy(..) )
import Data.Maybe (fromMaybe)
import Data.Functor.Compose (Compose (..))

import Data.Foldable(toList)
import Data.Traversable (mapAccumL)
import Data.Coerce (coerce)

import qualified Data.Map.Lazy as Map
import qualified Data.Map.Merge.Lazy as Map
import qualified Data.Vector as V

import Data.PTraversable

import Data.FunctorShape

import Control.Natural

-- | Data structure which represents partial natural transformations,
--   like usual 'Map' which represents partial functions.
--
-- @'Map' k v@ can be seen as a partial function.
--
-- @
-- m :: Map k v
-- (\k -> Data.Map.lookup k m) :: k -> Maybe v
-- @
-- 
-- Analogically, a @'NatMap' f g@ can be seen as a partial natural transformation,
-- in the same sense @Map@ represents a partial function.
-- 
-- @
-- nm :: NatMap f g
-- (\fa -> 'lookup' fa nm) :: forall a. f a -> Maybe (g a)
-- @
-- 
newtype NatMap (f :: Type -> Type) (g :: Type -> Type)
  = Mk (Map.Map (Shape f) (g Var))
type role NatMap nominal representational

deriving instance (WeakEq f, Eq (g Var)) => Eq (NatMap f g)
deriving instance (WeakOrd f, Ord (g Var)) => Ord (NatMap f g)

-- * Entry

data Entry f g = Entry (Shape f) (g Var)

deriving instance (Show (f Ignored), Show (g Var)) => Show (Entry f g)
deriving instance (WeakEq f, Eq (g Var)) => Eq (Entry f g)
deriving instance (WeakOrd f, Ord (g Var)) => Ord (Entry f g)

getKeyValue :: Entry f g -> (Shape f, g Var)
getKeyValue :: forall (f :: * -> *) (g :: * -> *). Entry f g -> (Shape f, g Var)
getKeyValue (Entry Shape f
fx g Var
gx) = (Shape f
fx, g Var
gx)

makeEntry :: (Traversable f, Traversable g, Ord k) => f k -> g k -> Maybe (Entry f g)
makeEntry :: forall (f :: * -> *) (g :: * -> *) k.
(Traversable f, Traversable g, Ord k) =>
f k -> g k -> Maybe (Entry f g)
makeEntry f k
fk g k
gk =
  do Map k Var
table <- [k] -> Maybe (Map k Var)
forall k. Ord k => [k] -> Maybe (Map k Var)
makeKeyTable (f k -> [k]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f k
fk)
     g Var
gx <- (k -> Maybe Var) -> g k -> Maybe (g Var)
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) -> g a -> f (g b)
traverse (\k
k -> k -> Map k Var -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k Var
table) g k
gk
     Entry f g -> Maybe (Entry f g)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape f -> g Var -> Entry f g
forall (f :: * -> *) (g :: * -> *). Shape f -> g Var -> Entry f g
Entry (f k -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f k
fk) g Var
gx)

makeIdEntry :: (Traversable f) => f k -> Entry f f
makeIdEntry :: forall (f :: * -> *) k. Traversable f => f k -> Entry f f
makeIdEntry f k
fk = Shape f -> f Var -> Entry f f
forall (f :: * -> *) (g :: * -> *). Shape f -> g Var -> Entry f g
Entry (f k -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f k
fk) (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices (f k -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f k
fk))

makeKeyTable :: Ord k => [k] -> Maybe (Map.Map k Var)
makeKeyTable :: forall k. Ord k => [k] -> Maybe (Map k Var)
makeKeyTable [k]
ks
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map k Var -> Int
forall k a. Map k a -> Int
Map.size Map k Var
table = Map k Var -> Maybe (Map k Var)
forall a. a -> Maybe a
Just Map k Var
table
  | Bool
otherwise = Maybe (Map k Var)
forall a. Maybe a
Nothing
  where
    n :: Int
n = [k] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [k]
ks
    table :: Map k Var
table = [(k, Var)] -> Map k Var
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([k] -> [Var] -> [(k, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [k]
ks (Int -> Var
Var (Int -> Var) -> [Int] -> [Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0..]))

unsafeMakeEntry :: (Traversable f, Functor g, Ord k) => f k -> g k -> Entry f g
unsafeMakeEntry :: forall (f :: * -> *) (g :: * -> *) k.
(Traversable f, Functor g, Ord k) =>
f k -> g k -> Entry f g
unsafeMakeEntry f k
fk g k
gk =
  case [k] -> Maybe (Map k Var)
forall k. Ord k => [k] -> Maybe (Map k Var)
makeKeyTable (f k -> [k]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f k
fk) of
    Just Map k Var
table -> Shape f -> g Var -> Entry f g
forall (f :: * -> *) (g :: * -> *). Shape f -> g Var -> Entry f g
Entry (f k -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f k
fk) ((Map k Var
table Map k Var -> k -> Var
forall k a. Ord k => Map k a -> k -> a
Map.!) (k -> Var) -> g k -> g Var
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g k
gk)
    Maybe (Map k Var)
Nothing -> String -> Entry f g
forall a. HasCallStack => String -> a
error String
"bad LHS functor!"

-- * Construction

empty :: NatMap f g
empty :: forall (f :: * -> *) (g :: * -> *). NatMap f g
empty = Map (Shape f) (g Var) -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk Map (Shape f) (g Var)
forall k a. Map k a
Map.empty

singleton :: Entry f g -> NatMap f g
singleton :: forall (f :: * -> *) (g :: * -> *). Entry f g -> NatMap f g
singleton (Entry Shape f
f1 g Var
gx) = Map (Shape f) (g Var) -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (g Var) -> NatMap f g)
-> Map (Shape f) (g Var) -> NatMap f g
forall a b. (a -> b) -> a -> b
$ Shape f -> g Var -> Map (Shape f) (g Var)
forall k a. k -> a -> Map k a
Map.singleton Shape f
f1 g Var
gx

partialIdentity :: (WeakOrd f, Traversable f) => [f a] -> NatMap f f
partialIdentity :: forall (f :: * -> *) a.
(WeakOrd f, Traversable f) =>
[f a] -> NatMap f f
partialIdentity [f a]
fShapes = Map (Shape f) (f Var) -> NatMap f f
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (f Var) -> NatMap f f)
-> Map (Shape f) (f Var) -> NatMap f f
forall a b. (a -> b) -> a -> b
$ [(Shape f, f Var)] -> Map (Shape f) (f Var)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
f1, Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
f1)) | f a
f1 <- [f a]
fShapes ]

fromEntries :: (WeakOrd f) => [Entry f g] -> NatMap f g
fromEntries :: forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
[Entry f g] -> NatMap f g
fromEntries [Entry f g]
es = Map (Shape f) (g Var) -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (g Var) -> NatMap f g)
-> Map (Shape f) (g Var) -> NatMap f g
forall a b. (a -> b) -> a -> b
$ [(Shape f, g Var)] -> Map (Shape f) (g Var)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (Entry f g -> (Shape f, g Var)
forall (f :: * -> *) (g :: * -> *). Entry f g -> (Shape f, g Var)
getKeyValue (Entry f g -> (Shape f, g Var))
-> [Entry f g] -> [(Shape f, g Var)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Entry f g]
es)

insert :: WeakOrd f => Entry f g -> NatMap f g -> NatMap f g
insert :: forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
Entry f g -> NatMap f g -> NatMap f g
insert (Entry Shape f
f1 g Var
gx) (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (g Var) -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (g Var) -> NatMap f g)
-> Map (Shape f) (g Var) -> NatMap f g
forall a b. (a -> b) -> a -> b
$ Shape f -> g Var -> Map (Shape f) (g Var) -> Map (Shape f) (g Var)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Shape f
f1 g Var
gx Map (Shape f) (g Var)
m

delete :: WeakOrd f => f any -> NatMap f g -> NatMap f g
delete :: forall (f :: * -> *) any (g :: * -> *).
WeakOrd f =>
f any -> NatMap f g -> NatMap f g
delete f any
fx (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (g Var) -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (g Var) -> NatMap f g)
-> Map (Shape f) (g Var) -> NatMap f g
forall a b. (a -> b) -> a -> b
$ Shape f -> Map (Shape f) (g Var) -> Map (Shape f) (g Var)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (f any -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f any
fx) Map (Shape f) (g Var)
m

-- * Queries

size :: NatMap f g -> Int
size :: forall (f :: * -> *) (g :: * -> *). NatMap f g -> Int
size (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (g Var) -> Int
forall k a. Map k a -> Int
Map.size Map (Shape f) (g Var)
m

member, notMember :: WeakOrd f => f a -> NatMap f g -> Bool
member :: forall (f :: * -> *) a (g :: * -> *).
WeakOrd f =>
f a -> NatMap f g -> Bool
member f a
fa (Mk Map (Shape f) (g Var)
m) = Shape f -> Map (Shape f) (g Var) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (g Var)
m
notMember :: forall (f :: * -> *) a (g :: * -> *).
WeakOrd f =>
f a -> NatMap f g -> Bool
notMember f a
fa = Bool -> Bool
not (Bool -> Bool) -> (NatMap f g -> Bool) -> NatMap f g -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> NatMap f g -> Bool
forall (f :: * -> *) a (g :: * -> *).
WeakOrd f =>
f a -> NatMap f g -> Bool
member f a
fa

lookup :: (WeakOrd f, Foldable f, Functor g) => f a -> NatMap f g -> Maybe (g a)
lookup :: forall (f :: * -> *) (g :: * -> *) a.
(WeakOrd f, Foldable f, Functor g) =>
f a -> NatMap f g -> Maybe (g a)
lookup f a
fa (Mk Map (Shape f) (g Var)
m) = (Var -> a) -> g Var -> g a
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
as (Int -> a) -> (Var -> Int) -> Var -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var -> Int
unVar) (g Var -> g a) -> Maybe (g Var) -> Maybe (g a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Shape f -> Map (Shape f) (g Var) -> Maybe (g Var)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (g Var)
m
  where
    as :: Vector a
as = [a] -> Vector a
forall a. [a] -> Vector a
V.fromList (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
fa)

lookup_ :: (WeakOrd f, Functor g) => Shape f -> NatMap f g -> Maybe (g Var)
lookup_ :: forall (f :: * -> *) (g :: * -> *).
(WeakOrd f, Functor g) =>
Shape f -> NatMap f g -> Maybe (g Var)
lookup_ Shape f
f1 (Mk Map (Shape f) (g Var)
m) = Shape f -> Map (Shape f) (g Var) -> Maybe (g Var)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Shape f
f1 Map (Shape f) (g Var)
m

keys :: NatMap f g -> [Shape f]
keys :: forall (f :: * -> *) (g :: * -> *). NatMap f g -> [Shape f]
keys (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (g Var) -> [Shape f]
forall k a. Map k a -> [k]
Map.keys Map (Shape f) (g Var)
m

toEntries :: Traversable f => NatMap f g -> [Entry f g]
toEntries :: forall (f :: * -> *) (g :: * -> *).
Traversable f =>
NatMap f g -> [Entry f g]
toEntries (Mk Map (Shape f) (g Var)
m) = [ Shape f -> g Var -> Entry f g
forall (f :: * -> *) (g :: * -> *). Shape f -> g Var -> Entry f g
Entry Shape f
f1 g Var
gx | (Shape f
f1, g Var
gx) <- Map (Shape f) (g Var) -> [(Shape f, g Var)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map (Shape f) (g Var)
m ]

-- keysSet :: NatMap f g -> ShapeSet f

-- * Mapping, Filtering, Traversing

map1 :: (forall a. g a -> h a) -> NatMap f g -> NatMap f h
map1 :: forall (g :: * -> *) (h :: * -> *) (f :: * -> *).
(forall a. g a -> h a) -> NatMap f g -> NatMap f h
map1 forall a. g a -> h a
gh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk ((g Var -> h Var) -> Map (Shape f) (g Var) -> Map (Shape f) (h Var)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map g Var -> h Var
forall a. g a -> h a
gh Map (Shape f) (g Var)
m)

mapMaybe1 :: (forall a. g a -> Maybe (h a)) -> NatMap f g -> NatMap f h
mapMaybe1 :: forall (g :: * -> *) (h :: * -> *) (f :: * -> *).
(forall a. g a -> Maybe (h a)) -> NatMap f g -> NatMap f h
mapMaybe1 forall a. g a -> Maybe (h a)
gh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk ((g Var -> Maybe (h Var))
-> Map (Shape f) (g Var) -> Map (Shape f) (h Var)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe g Var -> Maybe (h Var)
forall a. g a -> Maybe (h a)
gh Map (Shape f) (g Var)
m)

traverse1 :: Applicative m => (forall a. g a -> m (h a)) -> NatMap f g -> m (NatMap f h)
traverse1 :: forall (m :: * -> *) (g :: * -> *) (h :: * -> *) (f :: * -> *).
Applicative m =>
(forall a. g a -> m (h a)) -> NatMap f g -> m (NatMap f h)
traverse1 forall a. g a -> m (h a)
gh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> m (Map (Shape f) (h Var)) -> m (NatMap f h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (g Var -> m (h Var))
-> Map (Shape f) (g Var) -> m (Map (Shape f) (h Var))
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) -> Map (Shape f) a -> f (Map (Shape f) b)
traverse g Var -> m (h Var)
forall a. g a -> m (h a)
gh Map (Shape f) (g Var)
m

wither1 :: Applicative m => (forall a. g a -> m (Maybe (h a))) -> NatMap f g -> m (NatMap f h)
wither1 :: forall (m :: * -> *) (g :: * -> *) (h :: * -> *) (f :: * -> *).
Applicative m =>
(forall a. g a -> m (Maybe (h a))) -> NatMap f g -> m (NatMap f h)
wither1 forall a. g a -> m (Maybe (h a))
gh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> (Map (Shape f) (Maybe (h Var)) -> Map (Shape f) (h Var))
-> Map (Shape f) (Maybe (h Var))
-> NatMap f h
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe (h Var) -> Maybe (h Var))
-> Map (Shape f) (Maybe (h Var)) -> Map (Shape f) (h Var)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe Maybe (h Var) -> Maybe (h Var)
forall a. a -> a
id (Map (Shape f) (Maybe (h Var)) -> NatMap f h)
-> m (Map (Shape f) (Maybe (h Var))) -> m (NatMap f h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (g Var -> m (Maybe (h Var)))
-> Map (Shape f) (g Var) -> m (Map (Shape f) (Maybe (h Var)))
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) -> Map (Shape f) a -> f (Map (Shape f) b)
traverse g Var -> m (Maybe (h Var))
forall a. g a -> m (Maybe (h a))
gh Map (Shape f) (g Var)
m

mapWithKey1 :: Traversable f => (forall a. Ord a => f a -> g a -> h a) -> NatMap f g -> NatMap f h
mapWithKey1 :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
Traversable f =>
(forall a. Ord a => f a -> g a -> h a) -> NatMap f g -> NatMap f h
mapWithKey1 forall a. Ord a => f a -> g a -> h a
fgh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> Map (Shape f) (h Var) -> NatMap f h
forall a b. (a -> b) -> a -> b
$ (Shape f -> g Var -> h Var)
-> Map (Shape f) (g Var) -> Map (Shape f) (h Var)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\Shape f
f1 g Var
gx -> f Var -> g Var -> h Var
forall a. Ord a => f a -> g a -> h a
fgh (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1) g Var
gx) Map (Shape f) (g Var)
m

mapMaybeWithKey1 :: Traversable f => (forall a. Ord a => f a -> g a -> Maybe (h a)) -> NatMap f g -> NatMap f h
mapMaybeWithKey1 :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
Traversable f =>
(forall a. Ord a => f a -> g a -> Maybe (h a))
-> NatMap f g -> NatMap f h
mapMaybeWithKey1 forall a. Ord a => f a -> g a -> Maybe (h a)
fgh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> Map (Shape f) (h Var) -> NatMap f h
forall a b. (a -> b) -> a -> b
$ (Shape f -> g Var -> Maybe (h Var))
-> Map (Shape f) (g Var) -> Map (Shape f) (h Var)
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybeWithKey (\Shape f
f1 g Var
gx -> f Var -> g Var -> Maybe (h Var)
forall a. Ord a => f a -> g a -> Maybe (h a)
fgh (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1) g Var
gx) Map (Shape f) (g Var)
m

traverseWithKey1 :: (Traversable f, Applicative m) => (forall a. Ord a => f a -> g a -> m (h a)) -> NatMap f g -> m (NatMap f h)
traverseWithKey1 :: forall (f :: * -> *) (m :: * -> *) (g :: * -> *) (h :: * -> *).
(Traversable f, Applicative m) =>
(forall a. Ord a => f a -> g a -> m (h a))
-> NatMap f g -> m (NatMap f h)
traverseWithKey1 forall a. Ord a => f a -> g a -> m (h a)
fgh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> m (Map (Shape f) (h Var)) -> m (NatMap f h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Shape f -> g Var -> m (h Var))
-> Map (Shape f) (g Var) -> m (Map (Shape f) (h Var))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey (\Shape f
f1 g Var
gx -> f Var -> g Var -> m (h Var)
forall a. Ord a => f a -> g a -> m (h a)
fgh (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1) g Var
gx) Map (Shape f) (g Var)
m

witherWithKey1 :: (Traversable f, Applicative m) => (forall a. Ord a => f a -> g a -> m (Maybe (h a))) -> NatMap f g -> m (NatMap f h)
witherWithKey1 :: forall (f :: * -> *) (m :: * -> *) (g :: * -> *) (h :: * -> *).
(Traversable f, Applicative m) =>
(forall a. Ord a => f a -> g a -> m (Maybe (h a)))
-> NatMap f g -> m (NatMap f h)
witherWithKey1 forall a. Ord a => f a -> g a -> m (Maybe (h a))
fgh (Mk Map (Shape f) (g Var)
m) = Map (Shape f) (h Var) -> NatMap f h
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (h Var) -> NatMap f h)
-> (Map (Shape f) (Maybe (h Var)) -> Map (Shape f) (h Var))
-> Map (Shape f) (Maybe (h Var))
-> NatMap f h
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe (h Var) -> Maybe (h Var))
-> Map (Shape f) (Maybe (h Var)) -> Map (Shape f) (h Var)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe Maybe (h Var) -> Maybe (h Var)
forall a. a -> a
id (Map (Shape f) (Maybe (h Var)) -> NatMap f h)
-> m (Map (Shape f) (Maybe (h Var))) -> m (NatMap f h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Shape f -> g Var -> m (Maybe (h Var)))
-> Map (Shape f) (g Var) -> m (Map (Shape f) (Maybe (h Var)))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey (\Shape f
f1 g Var
gx -> f Var -> g Var -> m (Maybe (h Var))
forall a. Ord a => f a -> g a -> m (Maybe (h a))
fgh (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1) g Var
gx) Map (Shape f) (g Var)
m

-- * Combinators

union :: forall f g. (WeakOrd f) => NatMap f g -> NatMap f g -> NatMap f g
union :: forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
NatMap f g -> NatMap f g -> NatMap f g
union = (Map (Shape f) (g Var)
 -> Map (Shape f) (g Var) -> Map (Shape f) (g Var))
-> NatMap f g -> NatMap f g -> NatMap f g
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union @(Shape f) @(g Var))

unionWith :: forall f g. (WeakOrd f) => (forall a. g a -> g a -> g a) -> NatMap f g -> NatMap f g -> NatMap f g
unionWith :: forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
(forall a. g a -> g a -> g a)
-> NatMap f g -> NatMap f g -> NatMap f g
unionWith forall a. g a -> g a -> g a
op = (Map (Shape f) (g Var)
 -> Map (Shape f) (g Var) -> Map (Shape f) (g Var))
-> NatMap f g -> NatMap f g -> NatMap f g
forall a b. Coercible a b => a -> b
coerce ((g Var -> g Var -> g Var)
-> Map (Shape f) (g Var)
-> Map (Shape f) (g Var)
-> Map (Shape f) (g Var)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith (forall a. g a -> g a -> g a
op @Var))

consistentUnion :: forall f g. (WeakOrd f, Eq (g Var)) => NatMap f g -> NatMap f g -> Maybe (NatMap f g)
consistentUnion :: forall (f :: * -> *) (g :: * -> *).
(WeakOrd f, Eq (g Var)) =>
NatMap f g -> NatMap f g -> Maybe (NatMap f g)
consistentUnion = (Map (Shape f) (g Var)
 -> Map (Shape f) (g Var) -> Maybe (Map (Shape f) (g Var)))
-> NatMap f g -> NatMap f g -> Maybe (NatMap f g)
forall a b. Coercible a b => a -> b
coerce (forall k v. (Ord k, Eq v) => Map k v -> Map k v -> Maybe (Map k v)
consistentUnionImpl @(Shape f) @(g Var))

consistentUnionImpl :: forall k v. (Ord k, Eq v) => Map.Map k v -> Map.Map k v -> Maybe (Map.Map k v)
consistentUnionImpl :: forall k v. (Ord k, Eq v) => Map k v -> Map k v -> Maybe (Map k v)
consistentUnionImpl =
  WhenMissing Maybe k v v
-> WhenMissing Maybe k v v
-> WhenMatched Maybe k v v v
-> Map k v
-> Map k v
-> Maybe (Map k v)
forall (f :: * -> *) k a c b.
(Applicative f, Ord k) =>
WhenMissing f k a c
-> WhenMissing f k b c
-> WhenMatched f k a b c
-> Map k a
-> Map k b
-> f (Map k c)
Map.mergeA
    WhenMissing Maybe k v v
forall (f :: * -> *) k x. Applicative f => WhenMissing f k x x
Map.preserveMissing
    WhenMissing Maybe k v v
forall (f :: * -> *) k x. Applicative f => WhenMissing f k x x
Map.preserveMissing
    ((k -> v -> v -> Maybe v) -> WhenMatched Maybe k v v v
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> f z) -> WhenMatched f k x y z
Map.zipWithAMatched (\k
_ v
v1 v
v2 -> if v
v1 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v2 then v -> Maybe v
forall a. a -> Maybe a
Just v
v1 else Maybe v
forall a. Maybe a
Nothing))

-- * Combinators as partial natural transformations

identity :: PTraversable f => NatMap f f
identity :: forall (f :: * -> *). PTraversable f => NatMap f f
identity = Map (Shape f) (f Var) -> NatMap f f
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (g Var) -> NatMap f g
Mk (Map (Shape f) (f Var) -> NatMap f f)
-> Map (Shape f) (f Var) -> NatMap f f
forall a b. (a -> b) -> a -> b
$ [(Shape f, f Var)] -> Map (Shape f) (f Var)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (Shape f
f1, Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1) | Shape f
f1 <- f () -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape (f () -> Shape f) -> [f ()] -> [Shape f]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [()] -> [f ()]
forall (t :: * -> *) (f :: * -> *) a.
(PTraversable t, Alternative f) =>
f a -> f (t a)
enum1 [()] ]

compose :: (WeakOrd f, WeakOrd g, Foldable g, Functor h) => NatMap g h -> NatMap f g -> NatMap f h
compose :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
(WeakOrd f, WeakOrd g, Foldable g, Functor h) =>
NatMap g h -> NatMap f g -> NatMap f h
compose NatMap g h
nm1 NatMap f g
nm2 = (forall a. g a -> Maybe (h a)) -> NatMap f g -> NatMap f h
forall (g :: * -> *) (h :: * -> *) (f :: * -> *).
(forall a. g a -> Maybe (h a)) -> NatMap f g -> NatMap f h
mapMaybe1 (\g a
gx -> g a -> NatMap g h -> Maybe (h a)
forall (f :: * -> *) (g :: * -> *) a.
(WeakOrd f, Foldable f, Functor g) =>
f a -> NatMap f g -> Maybe (g a)
lookup g a
gx NatMap g h
nm1) NatMap f g
nm2

outerNat :: (Traversable f, Traversable g, PTraversable h, Ord (f (h Ignored))) => NatMap f g -> NatMap (Compose f h) (Compose g h)
outerNat :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
(Traversable f, Traversable g, PTraversable h,
 Ord (f (h Ignored))) =>
NatMap f g -> NatMap (Compose f h) (Compose g h)
outerNat NatMap f g
nm = [Entry (Compose f h) (Compose g h)]
-> NatMap (Compose f h) (Compose g h)
forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
[Entry f g] -> NatMap f g
fromEntries ([Entry (Compose f h) (Compose g h)]
 -> NatMap (Compose f h) (Compose g h))
-> [Entry (Compose f h) (Compose g h)]
-> NatMap (Compose f h) (Compose g h)
forall a b. (a -> b) -> a -> b
$
  do Entry (Shape f a
f_) g Var
gx <- NatMap f g -> [Entry f g]
forall (f :: * -> *) (g :: * -> *).
Traversable f =>
NatMap f g -> [Entry f g]
toEntries NatMap f g
nm
     f (h ())
fh_ <- (a -> [h ()]) -> f a -> [f (h ())]
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) -> f a -> f (f b)
traverse ([h ()] -> a -> [h ()]
forall a b. a -> b -> a
const [h ()]
hShapes) f a
f_
     let Compose f (h Var)
fhy = Shape (Compose f h) -> Compose f h Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices (Compose f h () -> Shape (Compose f h)
forall (f :: * -> *) a. f a -> Shape f
Shape (f (h ()) -> Compose f h ()
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose f (h ())
fh_))
         hyVec :: Vector (h Var)
hyVec = [h Var] -> Vector (h Var)
forall a. [a] -> Vector a
V.fromList (f (h Var) -> [h Var]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f (h Var)
fhy)
         ghy :: g (h Var)
ghy = (\Var
x -> Vector (h Var) -> Int -> h Var
forall a. Vector a -> Int -> a
V.unsafeIndex Vector (h Var)
hyVec (Var -> Int
unVar Var
x)) (Var -> h Var) -> g Var -> g (h Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g Var
gx
     Entry (Compose f h) (Compose g h)
-> [Entry (Compose f h) (Compose g h)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape (Compose f h)
-> Compose g h Var -> Entry (Compose f h) (Compose g h)
forall (f :: * -> *) (g :: * -> *). Shape f -> g Var -> Entry f g
Entry (Compose f h Var -> Shape (Compose f h)
forall (f :: * -> *) a. f a -> Shape f
Shape (f (h Var) -> Compose f h Var
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose f (h Var)
fhy)) (g (h Var) -> Compose g h Var
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose g (h Var)
ghy))
  where hShapes :: [h ()]
hShapes = [()] -> [h ()]
forall (t :: * -> *) (f :: * -> *) a.
(PTraversable t, Alternative f) =>
f a -> f (t a)
enum1 [()]

innerNat :: (Traversable f, Traversable g, PTraversable h, WeakOrd f) => NatMap f g -> NatMap (Compose h f) (Compose h g)
innerNat :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
(Traversable f, Traversable g, PTraversable h, WeakOrd f) =>
NatMap f g -> NatMap (Compose h f) (Compose h g)
innerNat NatMap f g
nm = [Entry (Compose h f) (Compose h g)]
-> NatMap (Compose h f) (Compose h g)
forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
[Entry f g] -> NatMap f g
fromEntries ([Entry (Compose h f) (Compose h g)]
 -> NatMap (Compose h f) (Compose h g))
-> [Entry (Compose h f) (Compose h g)]
-> NatMap (Compose h f) (Compose h g)
forall a b. (a -> b) -> a -> b
$
  do h Var
hx <- Shape h -> h Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices (Shape h -> h Var) -> (h () -> Shape h) -> h () -> h Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. h () -> Shape h
forall (f :: * -> *) a. f a -> Shape f
Shape (h () -> h Var) -> [h ()] -> [h Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [()] -> [h ()]
forall (t :: * -> *) (f :: * -> *) a.
(PTraversable t, Alternative f) =>
f a -> f (t a)
enum1 [()]
     h (f (Var, Var), g (Var, Var))
hEntry <- (Var -> [(f (Var, Var), g (Var, Var))])
-> h Var -> [h (f (Var, Var), g (Var, Var))]
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) -> h a -> f (h b)
traverse (\Var
x -> Var -> (f Var, g Var) -> (f (Var, Var), g (Var, Var))
forall {f :: * -> *} {f :: * -> *} {a} {a} {a}.
(Functor f, Functor f) =>
a -> (f a, f a) -> (f (a, a), f (a, a))
appendVar Var
x ((f Var, g Var) -> (f (Var, Var), g (Var, Var)))
-> [(f Var, g Var)] -> [(f (Var, Var), g (Var, Var))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(f Var, g Var)]
es) h Var
hx
     let hf_xy :: h (f (Var, Var))
hf_xy = (f (Var, Var), g (Var, Var)) -> f (Var, Var)
forall a b. (a, b) -> a
fst ((f (Var, Var), g (Var, Var)) -> f (Var, Var))
-> h (f (Var, Var), g (Var, Var)) -> h (f (Var, Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> h (f (Var, Var), g (Var, Var))
hEntry
         hg_xy :: h (g (Var, Var))
hg_xy = (f (Var, Var), g (Var, Var)) -> g (Var, Var)
forall a b. (a, b) -> b
snd ((f (Var, Var), g (Var, Var)) -> g (Var, Var))
-> h (f (Var, Var), g (Var, Var)) -> h (g (Var, Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> h (f (Var, Var), g (Var, Var))
hEntry
     Entry (Compose h f) (Compose h g)
-> [Entry (Compose h f) (Compose h g)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Entry (Compose h f) (Compose h g)
 -> [Entry (Compose h f) (Compose h g)])
-> Entry (Compose h f) (Compose h g)
-> [Entry (Compose h f) (Compose h g)]
forall a b. (a -> b) -> a -> b
$ Compose h f (Var, Var)
-> Compose h g (Var, Var) -> Entry (Compose h f) (Compose h g)
forall (f :: * -> *) (g :: * -> *) k.
(Traversable f, Functor g, Ord k) =>
f k -> g k -> Entry f g
unsafeMakeEntry (h (f (Var, Var)) -> Compose h f (Var, Var)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose h (f (Var, Var))
hf_xy) (h (g (Var, Var)) -> Compose h g (Var, Var)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose h (g (Var, Var))
hg_xy)
  where
    es :: [(f Var, g Var)]
es = [ (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1, g Var
gy) | Entry Shape f
f1 g Var
gy <- NatMap f g -> [Entry f g]
forall (f :: * -> *) (g :: * -> *).
Traversable f =>
NatMap f g -> [Entry f g]
toEntries NatMap f g
nm ]
    appendVar :: a -> (f a, f a) -> (f (a, a), f (a, a))
appendVar a
x (f a
fy, f a
gy) = ((,) a
x (a -> (a, a)) -> f a -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fy, (,) a
x (a -> (a, a)) -> f a -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
gy)

horizontalCompose :: (Traversable f, Functor g, Traversable h, Functor j, Ord (f (h Ignored)))
  => NatMap f g -> NatMap h j -> NatMap (Compose f h) (Compose g j)
horizontalCompose :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *) (j :: * -> *).
(Traversable f, Functor g, Traversable h, Functor j,
 Ord (f (h Ignored))) =>
NatMap f g -> NatMap h j -> NatMap (Compose f h) (Compose g j)
horizontalCompose NatMap f g
fgMap NatMap h j
hjMap = [Entry (Compose f h) (Compose g j)]
-> NatMap (Compose f h) (Compose g j)
forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
[Entry f g] -> NatMap f g
fromEntries ([Entry (Compose f h) (Compose g j)]
 -> NatMap (Compose f h) (Compose g j))
-> [Entry (Compose f h) (Compose g j)]
-> NatMap (Compose f h) (Compose g j)
forall a b. (a -> b) -> a -> b
$
  do (f Var
fx, g Var
gx) <- [(f Var, g Var)]
fgEntries
     f (h (Var, Var), j (Var, Var))
fPairs <- (Var -> [(h (Var, Var), j (Var, Var))])
-> f Var -> [f (h (Var, Var), j (Var, Var))]
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) -> f a -> f (f b)
traverse (\Var
x -> Var -> (h Var, j Var) -> (h (Var, Var), j (Var, Var))
forall {f :: * -> *} {f :: * -> *} {a} {a} {a}.
(Functor f, Functor f) =>
a -> (f a, f a) -> (f (a, a), f (a, a))
appendVar Var
x ((h Var, j Var) -> (h (Var, Var), j (Var, Var)))
-> [(h Var, j Var)] -> [(h (Var, Var), j (Var, Var))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(h Var, j Var)]
hjEntries) f Var
fx
     let fh_xy :: f (h (Var, Var))
fh_xy = (h (Var, Var), j (Var, Var)) -> h (Var, Var)
forall a b. (a, b) -> a
fst ((h (Var, Var), j (Var, Var)) -> h (Var, Var))
-> f (h (Var, Var), j (Var, Var)) -> f (h (Var, Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (h (Var, Var), j (Var, Var))
fPairs
         fj_xy :: f (j (Var, Var))
fj_xy = (h (Var, Var), j (Var, Var)) -> j (Var, Var)
forall a b. (a, b) -> b
snd ((h (Var, Var), j (Var, Var)) -> j (Var, Var))
-> f (h (Var, Var), j (Var, Var)) -> f (j (Var, Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (h (Var, Var), j (Var, Var))
fPairs
         vecj_xy :: Vector (j (Var, Var))
vecj_xy = [j (Var, Var)] -> Vector (j (Var, Var))
forall a. [a] -> Vector a
V.fromList (f (j (Var, Var)) -> [j (Var, Var)]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f (j (Var, Var))
fj_xy)
         gj_xy :: g (j (Var, Var))
gj_xy = (\Var
x -> Vector (j (Var, Var)) -> Int -> j (Var, Var)
forall a. Vector a -> Int -> a
V.unsafeIndex Vector (j (Var, Var))
vecj_xy (Var -> Int
unVar Var
x)) (Var -> j (Var, Var)) -> g Var -> g (j (Var, Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g Var
gx
     Entry (Compose f h) (Compose g j)
-> [Entry (Compose f h) (Compose g j)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Entry (Compose f h) (Compose g j)
 -> [Entry (Compose f h) (Compose g j)])
-> Entry (Compose f h) (Compose g j)
-> [Entry (Compose f h) (Compose g j)]
forall a b. (a -> b) -> a -> b
$ Compose f h (Var, Var)
-> Compose g j (Var, Var) -> Entry (Compose f h) (Compose g j)
forall (f :: * -> *) (g :: * -> *) k.
(Traversable f, Functor g, Ord k) =>
f k -> g k -> Entry f g
unsafeMakeEntry (f (h (Var, Var)) -> Compose f h (Var, Var)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose f (h (Var, Var))
fh_xy) (g (j (Var, Var)) -> Compose g j (Var, Var)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose g (j (Var, Var))
gj_xy)
  where
    toPairs :: NatMap f g -> [(f Var, g Var)]
toPairs NatMap f g
nm = (Entry f g -> (f Var, g Var)) -> [Entry f g] -> [(f Var, g Var)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Entry Shape f
f1 g Var
gx) -> (Shape f -> f Var
forall (f :: * -> *). Traversable f => Shape f -> f Var
indices Shape f
f1, g Var
gx)) (NatMap f g -> [Entry f g]
forall (f :: * -> *) (g :: * -> *).
Traversable f =>
NatMap f g -> [Entry f g]
toEntries NatMap f g
nm)
    fgEntries :: [(f Var, g Var)]
fgEntries = NatMap f g -> [(f Var, g Var)]
forall {f :: * -> *} {g :: * -> *}.
Traversable f =>
NatMap f g -> [(f Var, g Var)]
toPairs NatMap f g
fgMap
    hjEntries :: [(h Var, j Var)]
hjEntries = NatMap h j -> [(h Var, j Var)]
forall {f :: * -> *} {g :: * -> *}.
Traversable f =>
NatMap f g -> [(f Var, g Var)]
toPairs NatMap h j
hjMap
    appendVar :: a -> (f a, f a) -> (f (a, a), f (a, a))
appendVar a
x (f a
fy, f a
gy) = ((,) a
x (a -> (a, a)) -> f a -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fy, (,) a
x (a -> (a, a)) -> f a -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
gy)

-- * Totalities

fullSize :: forall f g. (PTraversable f) => NatMap f g -> Int
fullSize :: forall (f :: * -> *) (g :: * -> *).
PTraversable f =>
NatMap f g -> Int
fullSize NatMap f g
_ = forall (t :: * -> *) (proxy :: (* -> *) -> *).
PTraversable t =>
proxy t -> Int -> Int
cardinality1 @f Proxy f
forall {k} (t :: k). Proxy t
Proxy Int
1

isTotal :: (PTraversable f) => NatMap f g -> Bool
isTotal :: forall (f :: * -> *) (g :: * -> *).
PTraversable f =>
NatMap f g -> Bool
isTotal NatMap f g
nm = NatMap f g -> Int
forall (f :: * -> *) (g :: * -> *). NatMap f g -> Int
size NatMap f g
nm Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== NatMap f g -> Int
forall (f :: * -> *) (g :: * -> *).
PTraversable f =>
NatMap f g -> Int
fullSize NatMap f g
nm

toTotal :: (PTraversable f, Functor g) => NatMap f g -> Maybe (f :~> g)
toTotal :: forall (f :: * -> *) (g :: * -> *).
(PTraversable f, Functor g) =>
NatMap f g -> Maybe (f :~> g)
toTotal NatMap f g
nm
  | NatMap f g -> Bool
forall (f :: * -> *) (g :: * -> *).
PTraversable f =>
NatMap f g -> Bool
isTotal NatMap f g
nm = (f :~> g) -> Maybe (f :~> g)
forall a. a -> Maybe a
Just ((f ~> g) -> f :~> g
forall {k} (f :: k -> *) (g :: k -> *). (f ~> g) -> f :~> g
NT (\f x
fx -> g x -> Maybe (g x) -> g x
forall a. a -> Maybe a -> a
fromMaybe g x
forall {a}. a
err (f x -> NatMap f g -> Maybe (g x)
forall (f :: * -> *) (g :: * -> *) a.
(WeakOrd f, Foldable f, Functor g) =>
f a -> NatMap f g -> Maybe (g a)
lookup f x
fx NatMap f g
nm)))
  | Bool
otherwise  = Maybe (f :~> g)
forall a. Maybe a
Nothing
  where
    err :: a
err = String -> a
forall a. HasCallStack => String -> a
error String
"isTotal but there's a missed key?"

-- * Utility

indices :: Traversable f => Shape f -> f Var
indices :: forall (f :: * -> *). Traversable f => Shape f -> f Var
indices (Shape f a
fk) = (Int, f Var) -> f Var
forall a b. (a, b) -> b
snd ((Int, f Var) -> f Var) -> (f a -> (Int, f Var)) -> f a -> f Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> a -> (Int, Var)) -> Int -> f a -> (Int, f Var)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
n a
_ -> (Int -> Int
forall a. Enum a => a -> a
succ Int
n, Int -> Var
Var Int
n)) Int
0 (f a -> f Var) -> f a -> f Var
forall a b. (a -> b) -> a -> b
$ f a
fk

-- | An opaque type representing syntactic variable.
newtype Var = Var { Var -> Int
unVar :: Int }
  deriving (Var -> Var -> Bool
(Var -> Var -> Bool) -> (Var -> Var -> Bool) -> Eq Var
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Var -> Var -> Bool
== :: Var -> Var -> Bool
$c/= :: Var -> Var -> Bool
/= :: Var -> Var -> Bool
Eq, Eq Var
Eq Var =>
(Var -> Var -> Ordering)
-> (Var -> Var -> Bool)
-> (Var -> Var -> Bool)
-> (Var -> Var -> Bool)
-> (Var -> Var -> Bool)
-> (Var -> Var -> Var)
-> (Var -> Var -> Var)
-> Ord Var
Var -> Var -> Bool
Var -> Var -> Ordering
Var -> Var -> Var
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
$ccompare :: Var -> Var -> Ordering
compare :: Var -> Var -> Ordering
$c< :: Var -> Var -> Bool
< :: Var -> Var -> Bool
$c<= :: Var -> Var -> Bool
<= :: Var -> Var -> Bool
$c> :: Var -> Var -> Bool
> :: Var -> Var -> Bool
$c>= :: Var -> Var -> Bool
>= :: Var -> Var -> Bool
$cmax :: Var -> Var -> Var
max :: Var -> Var -> Var
$cmin :: Var -> Var -> Var
min :: Var -> Var -> Var
Ord)

instance Show Var where
  show :: Var -> String
show (Var Int
n) = String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n