{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

module Data.Record.Anon.Internal.Plugin.TC.TyConSubst (
    TyConSubst -- opaque
  , mkTyConSubst
  , splitTyConApp_upTo
  ) where

import Data.Bifunctor
import Data.Either (partitionEithers)
import Data.Foldable (toList, asum)
import Data.List.NonEmpty (NonEmpty(..))
import Data.Map (Map)

import qualified Data.Map as Map

import Data.Record.Anon.Internal.Plugin.TC.EquivClasses
import Data.Record.Anon.Internal.Plugin.TC.GhcTcPluginAPI hiding ((<>))

{-------------------------------------------------------------------------------
  The main type

  TODO: maybe this could be sped up with
  <https://hackage.haskell.org/package/union-find>?
-------------------------------------------------------------------------------}

-- | Substitution for recognizing 'TyCon' applications modulo equalities
--
-- During constraint solving the set of " given " constraints includes so-called
-- "canonical equalities": equalities of the form
--
-- > var ~ typ                  (CTyEqCan)
-- > var ~ TyCon arg1 .. argN   (CFunEqCan, the TyCon will be a type family)
--
-- The problem we want to solve is recognizing if some type τ is of the form
--
-- > TyCon arg1 arg2 .. argN   (0 <= N)
--
-- modulo those canonical equalities. We limit the scope of what we try to do:
--
-- o We are only interested in recognizing types of the form above
--   (as opposed to general parsing-modulo-equalities).
-- o We will only use the canonical equalities as-is: we will not attempt to
--   derive any additional equalities from them (i.e. if, say, we know that
--   @x ~ T1@ and @x ~ T2@, we will not attempt to use the fact that this means
--   that @T1 ~ T2@, nor any derived conclusions thereof). We /will/ however
--   try to apply the canononical equalities as often as is necessary (e.g.,
--   first applying @x ~ T y@, then applying @y ~ T2@).
--
-- We solve this problem by constructing a 'TyConSubst': a possibly
-- non-deterministic substitution mapping type variables to types of the form
-- above (that is, a type constructor applied to some arguments).
--
-- We detail the construction of this substitution below (see documentation of
-- 'Classified' and 'process'), but once we have this substitution, the
-- recognition problem becomes easy:
--
-- 1. Without loss of generality, let τ be of the form @t arg1 arg2 .. argN@
-- 2. If @t@ is a 'TyCon', we're done.
-- 3. Otherwise, if @t@ is a variable @x@, lookup @x@ in the substitution; if
--    there is one (or more) mappings for @x@, then we have successfully
--    recognized τ to be of the form above. There is no need to apply the
--    substitution repeatedly.
--
-- The substitution is non-deterministic because there might be multiple
-- matches. For example, if we have
--
-- > type family Foo where
-- >   Foo = Int
--
-- then we might well have equalities @x ~ Int, x ~ Foo@ in scope, and so a type
-- @x@ would match two different 'TyCon's. What we do know, however, is that if
-- τ matches both @t arg1 .. argN@ and @t' arg1' .. argM'@ (possibly @N /= M@),
-- then
--
-- > t arg1 .. argN ~ t' arg1' .. argM'
--
-- If @t == t'@, we can conclude that the arguments are equal only if @t@ is
-- injective.
data TyConSubst = TyConSubst {
      -- | Mapping from (canonical) variables to 'TyCon' applications
      TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))

      -- | Map each variable to the canonical representative
      --
      -- See 'Classified' for a detailed discussion of canonical variables.
    , TyConSubst -> Map TcTyVar TcTyVar
tyConSubstCanon :: Map TcTyVar TcTyVar
    }

{-------------------------------------------------------------------------------
  Basic functionality for working with 'TyConSubst'
-------------------------------------------------------------------------------}

-- | Empty substitution
--
-- The canonical variables map is established once when the initial substitution
-- is generated and not updated thereafter.
tyConSubstEmpty :: Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty :: Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty Map TcTyVar TcTyVar
canon = TyConSubst {
      tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap   = forall k a. Map k a
Map.empty
    , tyConSubstCanon :: Map TcTyVar TcTyVar
tyConSubstCanon = Map TcTyVar TcTyVar
canon
    }

-- | Lookup a variable in the substitution
tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstCanon :: Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
..} = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TcTyVar
var' Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap
  where
    var' :: TcTyVar
    var' :: TcTyVar
var' = forall a. Ord a => Map a a -> a -> a
canonicalize Map TcTyVar TcTyVar
tyConSubstCanon TcTyVar
var

-- | Extend substitution with new bindings
tyConSubstExtend ::
     [(TcTyVar, (TyCon, [Type]))]
  -> TyConSubst -> TyConSubst
tyConSubstExtend :: [(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
new subst :: TyConSubst
subst@TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstCanon :: Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
..} = TyConSubst
subst {
      tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Semigroup a => a -> a -> a
(<>)
                        (forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type]))
aux) [(TcTyVar, (TyCon, [Type]))]
new)
                        Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap
    }
  where
    aux :: TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type]))
    aux :: TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type]))
aux TcTyVar
var (TyCon, [Type])
s = (forall a. Ord a => Map a a -> a -> a
canonicalize Map TcTyVar TcTyVar
tyConSubstCanon TcTyVar
var, (TyCon, [Type])
s forall a. a -> [a] -> NonEmpty a
:| [])

{-------------------------------------------------------------------------------
  Classification
-------------------------------------------------------------------------------}

-- | Classified canonical equality constraints
--
-- The first step in the construction of the 'TyConSubst' is to classify the
-- available canonical equalities as one of three categories, defined below.
data Classified = Classified {
      -- | " Obviously " productive mappings
      --
      -- An equality @var := TyCon args@ is productive, because as soon as we
      -- apply it, we are done: we have successfully recognized a type as being
      -- an application of a concrete type constructor (note that we only ever
      -- apply the substitution to the head @t@ of a type @t args@, never to the
      -- arguments).
      Classified -> [(TcTyVar, (TyCon, [Type]))]
classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]

      -- | Extend equivalence class of variables
      --
      -- An equality @var1 := var2@ we will regard as extending the equivalence
      -- classes of variables (see 'constructEquivClasses').
    , Classified -> [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]

      -- | Substitutions we need to reconsider later
      --
      -- An equality @var1 := var2 args@ (with @args@ a non-empty list of
      -- arguments) is most problematic. Applying it /may/ allow us to make
      -- progress, but it may not (consider for example @var := var arg@). We
      -- will reconsider such equalities at the end (see 'process').
    , Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
    }

instance Semigroup Classified where
  Classified
c1 <> :: Classified -> Classified -> Classified
<> Classified
c2 = Classified {
        classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
classifiedProductive       = forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, (TyCon, [Type]))]
classifiedProductive
      , classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass = forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass
      , classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider       = forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider
      }
    where
      combine :: (Classified -> [a]) -> [a]
      combine :: forall a. (Classified -> [a]) -> [a]
combine Classified -> [a]
f = Classified -> [a]
f Classified
c1 forall a. [a] -> [a] -> [a]
++ Classified -> [a]
f Classified
c2

instance Monoid Classified where
  mempty :: Classified
mempty = [(TcTyVar, (TyCon, [Type]))]
-> [(TcTyVar, TcTyVar)]
-> [(TcTyVar, (TcTyVar, NonEmpty Type))]
-> Classified
Classified [] [] []

productive :: TcTyVar -> (TyCon, [Type]) -> Classified
productive :: TcTyVar -> (TyCon, [Type]) -> Classified
productive TcTyVar
var (TyCon
tyCon, [Type]
args) = forall a. Monoid a => a
mempty {
      classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
classifiedProductive = [(TcTyVar
var, (TyCon
tyCon, [Type]
args))]
    }

extendEquivClass :: TcTyVar -> TcTyVar -> Classified
extendEquivClass :: TcTyVar -> TcTyVar -> Classified
extendEquivClass TcTyVar
var TcTyVar
var' = forall a. Monoid a => a
mempty {
      classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass = [(TcTyVar
var, TcTyVar
var')]
    }

reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider TcTyVar
var (TcTyVar
var', NonEmpty Type
args) = forall a. Monoid a => a
mempty {
      classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider = [(TcTyVar
var, (TcTyVar
var', NonEmpty Type
args))]
    }

-- | Classify a set of given constraints
--
-- See 'Classified' for details.
classify :: [Ct] -> Classified
classify :: [Ct] -> Classified
classify = Classified -> [Ct] -> Classified
go forall a. Monoid a => a
mempty
  where
    go :: Classified -> [Ct] -> Classified
    go :: Classified -> [Ct] -> Classified
go Classified
acc []     = Classified
acc
    go Classified
acc (Ct
c:[Ct]
cs) =
        case Ct -> Maybe (TcTyVar, Type)
isCanonicalVarEq Ct
c of
          Just (TcTyVar
var, Type -> (Type, [Type])
splitAppTys -> (Type
fn, [Type]
args))
            | Just TyCon
tyCon <- Type -> Maybe TyCon
tyConAppTyCon_maybe Type
fn ->
                Classified -> [Ct] -> Classified
go (TcTyVar -> (TyCon, [Type]) -> Classified
productive TcTyVar
var (TyCon
tyCon, [Type]
args) forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
            | Just TcTyVar
var' <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn, forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
args ->
                Classified -> [Ct] -> Classified
go (TcTyVar -> TcTyVar -> Classified
extendEquivClass TcTyVar
var TcTyVar
var' forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
            | Just TcTyVar
var' <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn, Type
x:[Type]
xs <- [Type]
args ->
                Classified -> [Ct] -> Classified
go (TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider TcTyVar
var (TcTyVar
var', Type
x forall a. a -> [a] -> NonEmpty a
:| [Type]
xs) forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
          Maybe (TcTyVar, Type)
_otherwise ->
            Classified -> [Ct] -> Classified
go Classified
acc [Ct]
cs

{-------------------------------------------------------------------------------
  Processing
-------------------------------------------------------------------------------}

-- | Construct 'TyCon' substitution from classified equality constraints
--
-- The difficult part in constructing this substitution are the equalities of
-- the form @var1 ~ var2 args@, which we ear-marked as "to reconsider" during
-- classification.
--
-- We will do this iteratively:
--
-- o We first construct a set of variable equivalence classes based on
--   'classifiedExtendEquivClass' (using 'constructEquivClasses'), and use that
--   along with the "obviously productive" equalities ('classifiedProductive')
--   as the initial value of the accumulator (a 'TyConSubst').
-- o We then repeatedly consider the remaining equalities. Whenever there is
--   a substitution available in the accumulator for @var2@ which turns it into
--   a type of the form @TyCon args'@, we add @var1 := TyCon args' args@ to the
--   accumulator.
-- o We keep doing this until we can make no more progress.
--
-- The functions for working with 'TyConSubst' take the variable equivalence
-- classes into acocunt, so we do not need to do that here.
--
-- Two observations:
--
-- o This process must terminate: there are a finite number of constraints
--   to consider, and whenever we apply a substitution from the accumulator,
--   we get an "obviously productive" substitution: we do not create new work
--   in the loop.
-- o We may end up ignoring some substitutions: if there is a substitution
--   @var1 := var2 args@ and we don't have any (productive) substitutions for
--   @var2@, we will just ignore it.
--
-- A note on recursive bindings: a direct or indirect recursive binding
--
-- > x := x args1      x := y args1
-- >                   y := x args2
--
-- where @args1, args2@ are non-empty lists of arguments, /cannot/ be relevant:
-- if they were, that would imply that there is some type constructor (regular
-- datatype or type family) which can be applied to an arbitrary number of
-- arguments. Such datatypes or type families cannot be defined in Haskell.
-- We therefore take no special care in handling recursive bindings, other than
-- to note (as we did above) that the process must terminate.
process :: Classified -> TyConSubst
process :: Classified -> TyConSubst
process Classified{[(TcTyVar, (TyCon, [Type]))]
[(TcTyVar, (TcTyVar, NonEmpty Type))]
[(TcTyVar, TcTyVar)]
classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
classifiedReconsider :: Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedExtendEquivClass :: Classified -> [(TcTyVar, TcTyVar)]
classifiedProductive :: Classified -> [(TcTyVar, (TyCon, [Type]))]
..} =
    TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go TyConSubst
initSubst [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider
  where
    initSubst :: TyConSubst
    initSubst :: TyConSubst
initSubst =
          [(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
classifiedProductive
        forall a b. (a -> b) -> a -> b
$ Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty (forall a. Ord a => [(a, a)] -> Map a a
constructEquivClasses [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass)

    go :: TyConSubst
       -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
       -> TyConSubst
    go :: TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go TyConSubst
acc [(TcTyVar, (TcTyVar, NonEmpty Type))]
rs =
        let ([(TcTyVar, (TyCon, [Type]))]
prod, [(TcTyVar, (TcTyVar, NonEmpty Type))]
rest) = forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply (TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
makeProductive [(TcTyVar, (TcTyVar, NonEmpty Type))]
rs in
        if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(TcTyVar, (TyCon, [Type]))]
prod
          then TyConSubst
acc -- No other equations can be made productive
          else TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go ([(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
prod TyConSubst
acc) [(TcTyVar, (TcTyVar, NonEmpty Type))]
rest
      where
        makeProductive ::
             (TcTyVar, (TcTyVar, NonEmpty Type))
          -> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
        makeProductive :: (TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
makeProductive (TcTyVar
var, (TcTyVar
var', NonEmpty Type
args)) =
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
aux)) (TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var' TyConSubst
acc)
          where
            aux :: TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
            aux :: TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
aux TyCon
tyCon [Type]
args' = (TcTyVar
var, (TyCon
tyCon, ([Type]
args' forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty Type
args)))

-- | Construct 'TyConSubst'
--
-- This is the main function that builds the 'TyConSubst' from the set of
-- " given " constraints. The actual work is done by 'classify' and 'process'.
mkTyConSubst :: [Ct] -> TyConSubst
mkTyConSubst :: [Ct] -> TyConSubst
mkTyConSubst = Classified -> TyConSubst
process forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ct] -> Classified
classify

{-------------------------------------------------------------------------------
  Using
-------------------------------------------------------------------------------}

-- | Like 'splitTyConApp_maybe', but taking canonical equalities into account
--
-- See 'TyConSubst' for a detailed discussion.
splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type]))
splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type]))
splitTyConApp_upTo TyConSubst
subst Type
typ = forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum [
      -- Direct match
      do TyCon
tyCon <- Type -> Maybe TyCon
tyConAppTyCon_maybe Type
fn
         forall (m :: * -> *) a. Monad m => a -> m a
return ((TyCon
tyCon, [Type]
args) forall a. a -> [a] -> NonEmpty a
:| [])

      -- Indirect match
    , do TcTyVar
var <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn
         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a. [a] -> [a] -> [a]
++ [Type]
args))) forall a b. (a -> b) -> a -> b
$ TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var TyConSubst
subst
    ]
  where
    (Type
fn, [Type]
args) = Type -> (Type, [Type])
splitAppTys Type
typ

{-------------------------------------------------------------------------------
  Outputable
-------------------------------------------------------------------------------}

instance Outputable TyConSubst where
  ppr :: TyConSubst -> SDoc
ppr TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstCanon :: Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
..} = SDoc -> SDoc
parens forall a b. (a -> b) -> a -> b
$
          String -> SDoc
text String
"TyConSubst"
      SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap
      SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr Map TcTyVar TcTyVar
tyConSubstCanon

{-------------------------------------------------------------------------------
  Internal auxiliary
-------------------------------------------------------------------------------}

-- | Attempt to apply a non-deterministic function to a list of values
--
-- Returns the successful results as well as the inputs on which the function
-- failed.
tryApply :: forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply :: forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply a -> Maybe (NonEmpty b)
f = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [Either a b] -> ([a], [b])
partitionEithers forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map a -> Either (NonEmpty b) a
f'
  where
    f' :: a -> Either (NonEmpty b) a
    f' :: a -> Either (NonEmpty b) a
f' a
a = forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall a b. b -> Either a b
Right a
a) forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ a -> Maybe (NonEmpty b)
f a
a