{-# LANGUAGE BangPatterns, CPP, DerivingStrategies, DerivingVia         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving, MultiWayIf, StandaloneDeriving #-}
{-# LANGUAGE TypeOperators                                              #-}
module Control.Subcategory.Semialign
  ( CSemialign(..), CAlign(..),
    csalign, cpadZip, cpadZipWith
  ) where
import           Control.Applicative                  (ZipList)
import           Control.Monad                        (forM_)
import           Control.Monad.ST.Strict              (runST)
import           Control.Subcategory.Functor
import           Control.Subcategory.Wrapper.Internal
import           Data.Bifunctor                       (Bifunctor (bimap))
import           Data.Coerce
import           Data.Containers
import           Data.Functor.Compose                 (Compose (..))
import           Data.Functor.Identity                (Identity)
import qualified Data.Functor.Product                 as SOP
import           Data.Hashable                        (Hashable)
import           Data.HashMap.Strict                  (HashMap)
import           Data.IntMap.Strict                   (IntMap)
import qualified Data.IntSet                          as IS
import           Data.List.NonEmpty                   (NonEmpty)
import           Data.Map                             (Map)
import           Data.MonoTraversable
import qualified Data.Primitive.Array                 as A
import qualified Data.Primitive.PrimArray             as PA
import qualified Data.Primitive.SmallArray            as SA
import           Data.Proxy                           (Proxy)
import           Data.Semialign
import           Data.Semigroup                       (Option (..))
import           Data.Sequence                        (Seq)
import qualified Data.Sequences                       as MT
import           Data.These                           (These (..), fromThese,
                                                       mergeThese)
import           Data.Tree                            (Tree)
import qualified Data.Vector                          as V
import qualified Data.Vector.Primitive                as P
import qualified Data.Vector.Storable                 as S
import qualified Data.Vector.Unboxed                  as U
import           GHC.Generics                         ((:*:) (..), (:.:) (..))

class CFunctor f => CSemialign f where
  {-# MINIMAL calignWith #-}
  calignWith
    :: (Dom f a, Dom f b, Dom f c)
    => (These a b -> c) -> f a -> f b -> f c
  calign
    :: (Dom f a, Dom f b, Dom f (These a b))
    => f a -> f b -> f (These a b)
  {-# INLINE [1] calign #-}
  calign = (These a b -> These a b) -> f a -> f b -> f (These a b)
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> These a b
forall a. a -> a
id

instance Semialign f => CSemialign (WrapFunctor f) where
  calignWith :: (These a b -> c)
-> WrapFunctor f a -> WrapFunctor f b -> WrapFunctor f c
calignWith = (These a b -> c)
-> WrapFunctor f a -> WrapFunctor f b -> WrapFunctor f c
forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith
  {-# INLINE [1] calignWith #-}
  calign :: WrapFunctor f a -> WrapFunctor f b -> WrapFunctor f (These a b)
calign = WrapFunctor f a -> WrapFunctor f b -> WrapFunctor f (These a b)
forall (f :: * -> *) a b.
Semialign f =>
f a -> f b -> f (These a b)
align
  {-# INLINE [1] calign #-}

instance {-# OVERLAPPING #-}  CSemialign (WrapMono IS.IntSet) where
  calignWith :: (These a b -> c)
-> WrapMono IntSet a -> WrapMono IntSet b -> WrapMono IntSet c
calignWith These a b -> c
f = forall r.
(Coercible (WrapMono IntSet (Element IntSet)) IntSet => r) -> r
forall mono r.
(Coercible (WrapMono mono (Element mono)) mono => r) -> r
withMonoCoercible @IS.IntSet ((Coercible (WrapMono IntSet (Element IntSet)) IntSet =>
  WrapMono IntSet a -> WrapMono IntSet b -> WrapMono IntSet c)
 -> WrapMono IntSet a -> WrapMono IntSet b -> WrapMono IntSet c)
-> (Coercible (WrapMono IntSet (Element IntSet)) IntSet =>
    WrapMono IntSet a -> WrapMono IntSet b -> WrapMono IntSet c)
-> WrapMono IntSet a
-> WrapMono IntSet b
-> WrapMono IntSet c
forall a b. (a -> b) -> a -> b
$
    forall b.
Coercible (IntSet -> IntSet -> IntSet) b =>
(IntSet -> IntSet -> IntSet) -> b
coerce @(IS.IntSet -> IS.IntSet -> IS.IntSet) ((IntSet -> IntSet -> IntSet)
 -> WrapMono IntSet a -> WrapMono IntSet b -> WrapMono IntSet c)
-> (IntSet -> IntSet -> IntSet)
-> WrapMono IntSet a
-> WrapMono IntSet b
-> WrapMono IntSet c
forall a b. (a -> b) -> a -> b
$ \ IntSet
l IntSet
r ->
    let ints :: IntSet
ints = IntSet
l IntSet -> IntSet -> IntSet
`IS.intersection` IntSet
r
    in [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions
          [ (Key -> Key) -> IntSet -> IntSet
IS.map (These a b -> c
f (These a b -> c) -> (a -> These a b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> These a b
forall a b. a -> These a b
This) IntSet
l
          , (Key -> Key) -> IntSet -> IntSet
IS.map (These a b -> c
f (These a b -> c) -> (b -> These a b) -> b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> These a b
forall a b. b -> These a b
That) IntSet
r
          , (Key -> Key) -> IntSet -> IntSet
IS.map (\Key
x -> These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ Key -> Key -> These Key Key
forall a b. a -> b -> These a b
These Key
x Key
x) IntSet
ints
          ]
  {-# INLINE [1] calignWith #-}

class CSemialign f => CAlign f where
  cnil :: Dom f a => f a

instance Align f => CAlign (WrapFunctor f) where
  cnil :: WrapFunctor f a
cnil = f a -> WrapFunctor f a
forall (f :: * -> *) a. f a -> WrapFunctor f a
WrapFunctor f a
forall (f :: * -> *) a. Align f => f a
nil
  {-# INLINE [1] cnil #-}

deriving via WrapFunctor [] instance CSemialign []
deriving via WrapFunctor [] instance CAlign []
deriving via WrapFunctor Maybe instance CSemialign Maybe
deriving via WrapFunctor Maybe instance CAlign Maybe
#if MIN_VERSION_semialign(1,1,0)
deriving via WrapFunctor Option instance CSemialign Option
deriving via WrapFunctor Option instance CAlign Option
#else
deriving newtype instance CSemialign Option
deriving newtype instance CAlign Option
#endif

deriving via WrapFunctor ZipList instance CSemialign ZipList
deriving via WrapFunctor ZipList instance CAlign ZipList
deriving via WrapFunctor Identity instance CSemialign Identity
deriving via WrapFunctor NonEmpty instance CSemialign NonEmpty
deriving via WrapFunctor IntMap instance CSemialign IntMap
deriving via WrapFunctor IntMap instance CAlign IntMap
deriving via WrapFunctor Tree instance CSemialign Tree
deriving via WrapFunctor Seq instance CSemialign Seq
deriving via WrapFunctor Seq instance CAlign Seq
deriving via WrapFunctor V.Vector instance CSemialign V.Vector
deriving via WrapFunctor V.Vector instance CAlign V.Vector
deriving via WrapFunctor Proxy instance CSemialign Proxy
deriving via WrapFunctor Proxy instance CAlign Proxy
deriving via WrapFunctor (Map k) instance Ord k => CSemialign (Map k)
deriving via WrapFunctor (Map k) instance Ord k => CAlign (Map k)
deriving via WrapFunctor (HashMap k)
  instance (Eq k, Hashable k) => CSemialign (HashMap k)
deriving via WrapFunctor (HashMap k)
  instance (Eq k, Hashable k) => CAlign (HashMap k)
deriving via WrapFunctor ((->) s) instance CSemialign ((->) s)

instance (CSemialign f, CSemialign g) => CSemialign (SOP.Product f g) where
  calign :: Product f g a -> Product f g b -> Product f g (These a b)
calign (SOP.Pair f a
a g a
b) (SOP.Pair f b
c g b
d) = f (These a b) -> g (These a b) -> Product f g (These a b)
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
SOP.Pair (f a -> f b -> f (These a b)
forall (f :: * -> *) a b.
(CSemialign f, Dom f a, Dom f b, Dom f (These a b)) =>
f a -> f b -> f (These a b)
calign f a
a f b
c) (g a -> g b -> g (These a b)
forall (f :: * -> *) a b.
(CSemialign f, Dom f a, Dom f b, Dom f (These a b)) =>
f a -> f b -> f (These a b)
calign g a
b g b
d)
  {-# INLINE [1] calign #-}
  calignWith :: (These a b -> c) -> Product f g a -> Product f g b -> Product f g c
calignWith These a b -> c
f (SOP.Pair f a
a g a
b) (SOP.Pair f b
c g b
d) =
    f c -> g c -> Product f g c
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
SOP.Pair ((These a b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f f a
a f b
c) ((These a b -> c) -> g a -> g b -> g c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f g a
b g b
d)
  {-# INLINE [1] calignWith #-}

instance (CAlign f, CAlign g) => CAlign (SOP.Product f g) where
  cnil :: Product f g a
cnil = f a -> g a -> Product f g a
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
SOP.Pair f a
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil g a
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil
  {-# INLINE [1] cnil #-}

instance (CSemialign f, CSemialign g) => CSemialign (f :*: g) where
  calign :: (:*:) f g a -> (:*:) f g b -> (:*:) f g (These a b)
calign ((:*:) f a
a g a
b) ((:*:) f b
c g b
d) = f (These a b) -> g (These a b) -> (:*:) f g (These a b)
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) (f a -> f b -> f (These a b)
forall (f :: * -> *) a b.
(CSemialign f, Dom f a, Dom f b, Dom f (These a b)) =>
f a -> f b -> f (These a b)
calign f a
a f b
c) (g a -> g b -> g (These a b)
forall (f :: * -> *) a b.
(CSemialign f, Dom f a, Dom f b, Dom f (These a b)) =>
f a -> f b -> f (These a b)
calign g a
b g b
d)
  {-# INLINE [1] calign #-}
  calignWith :: (These a b -> c) -> (:*:) f g a -> (:*:) f g b -> (:*:) f g c
calignWith These a b -> c
f ((:*:) f a
a g a
b) ((:*:) f b
c g b
d) =
    f c -> g c -> (:*:) f g c
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) ((These a b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f f a
a f b
c) ((These a b -> c) -> g a -> g b -> g c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f g a
b g b
d)
  {-# INLINE [1] calignWith #-}

instance (CAlign f, CAlign g) => CAlign (f :*: g) where
  cnil :: (:*:) f g a
cnil = f a
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil f a -> g a -> (:*:) f g a
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: g a
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil
  {-# INLINE [1] cnil #-}

instance (CSemialign f, CSemialign g) => CSemialign (Compose f g) where
  calignWith :: (These a b -> c) -> Compose f g a -> Compose f g b -> Compose f g c
calignWith These a b -> c
f (Compose f (g a)
x) (Compose f (g b)
y) = f (g c) -> Compose f g c
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose ((These (g a) (g b) -> g c) -> f (g a) -> f (g b) -> f (g c)
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These (g a) (g b) -> g c
g f (g a)
x f (g b)
y)
    where
      g :: These (g a) (g b) -> g c
g (This g a
ga)     = (a -> c) -> g a -> g c
forall (f :: * -> *) a b.
(CFunctor f, Dom f a, Dom f b) =>
(a -> b) -> f a -> f b
cmap (These a b -> c
f (These a b -> c) -> (a -> These a b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> These a b
forall a b. a -> These a b
This) g a
ga
      g (That g b
gb)     = (b -> c) -> g b -> g c
forall (f :: * -> *) a b.
(CFunctor f, Dom f a, Dom f b) =>
(a -> b) -> f a -> f b
cmap (These a b -> c
f (These a b -> c) -> (b -> These a b) -> b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> These a b
forall a b. b -> These a b
That) g b
gb
      g (These g a
ga g b
gb) = (These a b -> c) -> g a -> g b -> g c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f g a
ga g b
gb
  {-# INLINE [1] calignWith #-}

instance (CAlign f, CSemialign g) => CAlign (Compose f g) where
  cnil :: Compose f g a
cnil = f (g a) -> Compose f g a
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose f (g a)
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil
  {-# INLINE [1] cnil #-}

instance (CSemialign f, CSemialign g) => CSemialign ((:.:) f g) where
  calignWith :: (These a b -> c) -> (:.:) f g a -> (:.:) f g b -> (:.:) f g c
calignWith These a b -> c
f (Comp1 f (g a)
x) (Comp1 f (g b)
y) = f (g c) -> (:.:) f g c
forall k2 k1 (f :: k2 -> *) (g :: k1 -> k2) (p :: k1).
f (g p) -> (:.:) f g p
Comp1 ((These (g a) (g b) -> g c) -> f (g a) -> f (g b) -> f (g c)
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These (g a) (g b) -> g c
g f (g a)
x f (g b)
y)
    where
      g :: These (g a) (g b) -> g c
g (This g a
ga)     = (a -> c) -> g a -> g c
forall (f :: * -> *) a b.
(CFunctor f, Dom f a, Dom f b) =>
(a -> b) -> f a -> f b
cmap (These a b -> c
f (These a b -> c) -> (a -> These a b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> These a b
forall a b. a -> These a b
This) g a
ga
      g (That g b
gb)     = (b -> c) -> g b -> g c
forall (f :: * -> *) a b.
(CFunctor f, Dom f a, Dom f b) =>
(a -> b) -> f a -> f b
cmap (These a b -> c
f (These a b -> c) -> (b -> These a b) -> b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> These a b
forall a b. b -> These a b
That) g b
gb
      g (These g a
ga g b
gb) = (These a b -> c) -> g a -> g b -> g c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith These a b -> c
f g a
ga g b
gb
  {-# INLINE [1] calignWith #-}

instance (CAlign f, CSemialign g) => CAlign ((:.:) f g) where
  cnil :: (:.:) f g a
cnil = f (g a) -> (:.:) f g a
forall k2 k1 (f :: k2 -> *) (g :: k1 -> k2) (p :: k1).
f (g p) -> (:.:) f g p
Comp1 f (g a)
forall (f :: * -> *) a. (CAlign f, Dom f a) => f a
cnil
  {-# INLINE [1] cnil #-}

instance CSemialign U.Vector where
  calignWith :: (These a b -> c) -> Vector a -> Vector b -> Vector c
calignWith = (These a b -> c) -> Vector a -> Vector b -> Vector c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(These a b -> c) -> v a -> v b -> v c
alignVectorWith
  {-# INLINE [1] calignWith #-}

instance CAlign U.Vector where
  cnil :: Vector a
cnil = Vector a
forall a. Unbox a => Vector a
U.empty
  {-# INLINE [1] cnil #-}

instance CSemialign S.Vector where
  calignWith :: (These a b -> c) -> Vector a -> Vector b -> Vector c
calignWith = (These a b -> c) -> Vector a -> Vector b -> Vector c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(These a b -> c) -> v a -> v b -> v c
alignVectorWith
  {-# INLINE [1] calignWith #-}

instance CAlign S.Vector where
  cnil :: Vector a
cnil = Vector a
forall a. Storable a => Vector a
S.empty
  {-# INLINE [1] cnil #-}

instance CSemialign P.Vector where
  calignWith :: (These a b -> c) -> Vector a -> Vector b -> Vector c
calignWith = (These a b -> c) -> Vector a -> Vector b -> Vector c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(These a b -> c) -> v a -> v b -> v c
alignVectorWith
  {-# INLINE [1] calignWith #-}

instance CAlign P.Vector where
  cnil :: Vector a
cnil = Vector a
forall a. Prim a => Vector a
P.empty
  {-# INLINE [1] cnil #-}

instance CSemialign SA.SmallArray where
  calignWith :: (These a b -> c) -> SmallArray a -> SmallArray b -> SmallArray c
calignWith These a b -> c
f SmallArray a
l SmallArray b
r = (forall s. ST s (SmallArray c)) -> SmallArray c
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray c)) -> SmallArray c)
-> (forall s. ST s (SmallArray c)) -> SmallArray c
forall a b. (a -> b) -> a -> b
$ do
    let !lenL :: Key
lenL = SmallArray a -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length SmallArray a
l
        !lenR :: Key
lenR = SmallArray b -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length SmallArray b
r
        (Bool
isLftShort, Key
thresh, Key
len)
          | Key
lenL Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
lenR = (Bool
True, Key
lenL, Key
lenR)
          | Bool
otherwise = (Bool
False, Key
lenR, Key
lenL)
    SmallMutableArray s c
sa <- Key -> c -> ST s (SmallMutableArray (PrimState (ST s)) c)
forall (m :: * -> *) a.
PrimMonad m =>
Key -> a -> m (SmallMutableArray (PrimState m) a)
SA.newSmallArray Key
len ([Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Uninitialised element")
    [Key] -> (Key -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Key
0..Key
lenKey -> Key -> Key
forall a. Num a => a -> a -> a
-Key
1] ((Key -> ST s ()) -> ST s ()) -> (Key -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Key
n ->
      if  | Key
n Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
len -> () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Key
n Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
thresh ->
            SmallMutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Key -> a -> m ()
SA.writeSmallArray SmallMutableArray s c
SmallMutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> b -> These a b
forall a b. a -> b -> These a b
These
              (SmallArray a -> Key -> a
forall a. SmallArray a -> Key -> a
SA.indexSmallArray SmallArray a
l Key
n)
              (SmallArray b -> Key -> b
forall a. SmallArray a -> Key -> a
SA.indexSmallArray SmallArray b
r Key
n)
          | Bool
isLftShort ->
            SmallMutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Key -> a -> m ()
SA.writeSmallArray SmallMutableArray s c
SmallMutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ b -> These a b
forall a b. b -> These a b
That (b -> These a b) -> b -> These a b
forall a b. (a -> b) -> a -> b
$ SmallArray b -> Key -> b
forall a. SmallArray a -> Key -> a
SA.indexSmallArray SmallArray b
r Key
n
          | Bool
otherwise ->
            SmallMutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Key -> a -> m ()
SA.writeSmallArray SmallMutableArray s c
SmallMutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> These a b
forall a b. a -> These a b
This (a -> These a b) -> a -> These a b
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Key -> a
forall a. SmallArray a -> Key -> a
SA.indexSmallArray SmallArray a
l Key
n
    SmallMutableArray (PrimState (ST s)) c -> ST s (SmallArray c)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
SA.unsafeFreezeSmallArray SmallMutableArray s c
SmallMutableArray (PrimState (ST s)) c
sa
  {-# INLINE [1] calignWith #-}

instance CAlign SA.SmallArray where
  cnil :: SmallArray a
cnil = Key -> [a] -> SmallArray a
forall a. Key -> [a] -> SmallArray a
SA.smallArrayFromListN Key
0 []
  {-# INLINE [1] cnil #-}

instance CSemialign A.Array where
  calignWith :: (These a b -> c) -> Array a -> Array b -> Array c
calignWith These a b -> c
f Array a
l Array b
r = (forall s. ST s (Array c)) -> Array c
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array c)) -> Array c)
-> (forall s. ST s (Array c)) -> Array c
forall a b. (a -> b) -> a -> b
$ do
    let !lenL :: Key
lenL = Array a -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length Array a
l
        !lenR :: Key
lenR = Array b -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length Array b
r
        (Bool
isLftShort, Key
thresh, Key
len)
          | Key
lenL Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
lenR = (Bool
True, Key
lenL, Key
lenR)
          | Bool
otherwise = (Bool
False, Key
lenR, Key
lenL)
    MutableArray s c
sa <- Key -> c -> ST s (MutableArray (PrimState (ST s)) c)
forall (m :: * -> *) a.
PrimMonad m =>
Key -> a -> m (MutableArray (PrimState m) a)
A.newArray Key
len ([Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Uninitialised element")
    [Key] -> (Key -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Key
0..Key
lenKey -> Key -> Key
forall a. Num a => a -> a -> a
-Key
1] ((Key -> ST s ()) -> ST s ()) -> (Key -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Key
n ->
      if  | Key
n Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
len -> () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Key
n Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
thresh ->
            MutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Key -> a -> m ()
A.writeArray MutableArray s c
MutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> b -> These a b
forall a b. a -> b -> These a b
These
              (Array a -> Key -> a
forall a. Array a -> Key -> a
A.indexArray Array a
l Key
n)
              (Array b -> Key -> b
forall a. Array a -> Key -> a
A.indexArray Array b
r Key
n)
          | Bool
isLftShort ->
            MutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Key -> a -> m ()
A.writeArray MutableArray s c
MutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ b -> These a b
forall a b. b -> These a b
That (b -> These a b) -> b -> These a b
forall a b. (a -> b) -> a -> b
$ Array b -> Key -> b
forall a. Array a -> Key -> a
A.indexArray Array b
r Key
n
          | Bool
otherwise ->
            MutableArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Key -> a -> m ()
A.writeArray MutableArray s c
MutableArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> These a b
forall a b. a -> These a b
This (a -> These a b) -> a -> These a b
forall a b. (a -> b) -> a -> b
$ Array a -> Key -> a
forall a. Array a -> Key -> a
A.indexArray Array a
l Key
n
    MutableArray (PrimState (ST s)) c -> ST s (Array c)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
A.unsafeFreezeArray MutableArray s c
MutableArray (PrimState (ST s)) c
sa
  {-# INLINE [1] calignWith #-}

instance CAlign A.Array where
  cnil :: Array a
cnil = Key -> [Item (Array a)] -> Array a
forall l. IsList l => Key -> [Item l] -> l
A.fromListN Key
0 []
  {-# INLINE [1] cnil #-}

instance CSemialign PA.PrimArray where
  calignWith :: (These a b -> c) -> PrimArray a -> PrimArray b -> PrimArray c
calignWith These a b -> c
f PrimArray a
l PrimArray b
r = (forall s. ST s (PrimArray c)) -> PrimArray c
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray c)) -> PrimArray c)
-> (forall s. ST s (PrimArray c)) -> PrimArray c
forall a b. (a -> b) -> a -> b
$ do
    let !lenL :: Key
lenL = PrimArray a -> Key
forall a. Prim a => PrimArray a -> Key
PA.sizeofPrimArray PrimArray a
l
        !lenR :: Key
lenR = PrimArray b -> Key
forall a. Prim a => PrimArray a -> Key
PA.sizeofPrimArray PrimArray b
r
        (Bool
isLftShort, Key
thresh, Key
len)
          | Key
lenL Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
lenR = (Bool
True, Key
lenL, Key
lenR)
          | Bool
otherwise = (Bool
False, Key
lenR, Key
lenL)
    MutablePrimArray s c
sa <- Key -> ST s (MutablePrimArray (PrimState (ST s)) c)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Key -> m (MutablePrimArray (PrimState m) a)
PA.newPrimArray Key
len
    [Key] -> (Key -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Key
0..Key
lenKey -> Key -> Key
forall a. Num a => a -> a -> a
-Key
1] ((Key -> ST s ()) -> ST s ()) -> (Key -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Key
n ->
      if  | Key
n Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
len -> () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Key
n Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
thresh ->
            MutablePrimArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Key -> a -> m ()
PA.writePrimArray MutablePrimArray s c
MutablePrimArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> b -> These a b
forall a b. a -> b -> These a b
These
              (PrimArray a -> Key -> a
forall a. Prim a => PrimArray a -> Key -> a
PA.indexPrimArray PrimArray a
l Key
n)
              (PrimArray b -> Key -> b
forall a. Prim a => PrimArray a -> Key -> a
PA.indexPrimArray PrimArray b
r Key
n)
          | Bool
isLftShort ->
            MutablePrimArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Key -> a -> m ()
PA.writePrimArray MutablePrimArray s c
MutablePrimArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ b -> These a b
forall a b. b -> These a b
That (b -> These a b) -> b -> These a b
forall a b. (a -> b) -> a -> b
$ PrimArray b -> Key -> b
forall a. Prim a => PrimArray a -> Key -> a
PA.indexPrimArray PrimArray b
r Key
n
          | Bool
otherwise ->
            MutablePrimArray (PrimState (ST s)) c -> Key -> c -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Key -> a -> m ()
PA.writePrimArray MutablePrimArray s c
MutablePrimArray (PrimState (ST s)) c
sa Key
n
            (c -> ST s ()) -> c -> ST s ()
forall a b. (a -> b) -> a -> b
$ These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> These a b
forall a b. a -> These a b
This (a -> These a b) -> a -> These a b
forall a b. (a -> b) -> a -> b
$ PrimArray a -> Key -> a
forall a. Prim a => PrimArray a -> Key -> a
PA.indexPrimArray PrimArray a
l Key
n
    MutablePrimArray (PrimState (ST s)) c -> ST s (PrimArray c)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
PA.unsafeFreezePrimArray MutablePrimArray s c
MutablePrimArray (PrimState (ST s)) c
sa
  {-# INLINE [1] calignWith #-}

instance CAlign PA.PrimArray where
  cnil :: PrimArray a
cnil = Key -> [a] -> PrimArray a
forall a. Prim a => Key -> [a] -> PrimArray a
PA.primArrayFromListN Key
0 []
  {-# INLINE [1] cnil #-}

instance (MT.IsSequence mono, MonoZip mono)
  => CSemialign (WrapMono mono) where
  calignWith :: (These a b -> c)
-> WrapMono mono a -> WrapMono mono b -> WrapMono mono c
calignWith These a b -> c
f = (mono -> mono -> mono)
-> WrapMono mono a -> WrapMono mono b -> WrapMono mono c
coerce mono -> mono -> mono
go
    where
      go :: mono -> mono -> mono
      go :: mono -> mono -> mono
go mono
ls mono
rs
        | Key
lenL Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
lenR = (Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
forall mono.
MonoZip mono =>
(Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
ozipWith ((These a b -> c) -> (a -> These a b) -> a -> c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap These a b -> c
f ((a -> These a b) -> a -> c)
-> (a -> a -> These a b) -> a -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> These a b
forall a b. a -> b -> These a b
These) mono
ls mono
rs
        | Key
lenL Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
lenR  =
            (Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
forall mono.
MonoZip mono =>
(Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
ozipWith ((These a b -> c) -> (a -> These a b) -> a -> c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap These a b -> c
f ((a -> These a b) -> a -> c)
-> (a -> a -> These a b) -> a -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> These a b
forall a b. a -> b -> These a b
These) mono
ls mono
rs
            mono -> mono -> mono
forall a. Semigroup a => a -> a -> a
<> (Element mono -> Element mono) -> mono -> mono
forall mono.
MonoFunctor mono =>
(Element mono -> Element mono) -> mono -> mono
omap (These a b -> c
f (These a b -> c) -> (b -> These a b) -> b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> These a b
forall a b. b -> These a b
That) (Index mono -> mono -> mono
forall seq. IsSequence seq => Index seq -> seq -> seq
MT.drop (Key -> Index mono
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
lenL) mono
rs)
        | Bool
otherwise  =
            (Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
forall mono.
MonoZip mono =>
(Element mono -> Element mono -> Element mono)
-> mono -> mono -> mono
ozipWith ((These a b -> c) -> (a -> These a b) -> a -> c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap These a b -> c
f ((a -> These a b) -> a -> c)
-> (a -> a -> These a b) -> a -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> These a b
forall a b. a -> b -> These a b
These) mono
ls mono
rs
            mono -> mono -> mono
forall a. Semigroup a => a -> a -> a
<> (Element mono -> Element mono) -> mono -> mono
forall mono.
MonoFunctor mono =>
(Element mono -> Element mono) -> mono -> mono
omap (These a b -> c
f (These a b -> c) -> (a -> These a b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> These a b
forall a b. a -> These a b
This) (Index mono -> mono -> mono
forall seq. IsSequence seq => Index seq -> seq -> seq
MT.drop (Key -> Index mono
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
lenL) mono
ls)
        where lenL :: Key
lenL = mono -> Key
forall mono. MonoFoldable mono => mono -> Key
olength mono
ls
              lenR :: Key
lenR = mono -> Key
forall mono. MonoFoldable mono => mono -> Key
olength mono
rs

instance (MT.IsSequence mono, MonoZip mono)
  => CAlign (WrapMono mono) where
  cnil :: WrapMono mono a
cnil = mono -> WrapMono mono a
forall b mono.
(b ~ Element mono, b ~ Element mono) =>
mono -> WrapMono mono b
WrapMono mono
forall a. Monoid a => a
mempty

csalign :: (CSemialign f, Dom f a, Semigroup a)
  => f a -> f a -> f a
{-# INLINE [1] csalign #-}
csalign :: f a -> f a -> f a
csalign = (These a a -> a) -> f a -> f a -> f a
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith ((These a a -> a) -> f a -> f a -> f a)
-> (These a a -> a) -> f a -> f a -> f a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> These a a -> a
forall a. (a -> a -> a) -> These a a -> a
mergeThese a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

cpadZip
  :: (CSemialign f, Dom f a, Dom f b, Dom f (Maybe a, Maybe b))
  => f a -> f b -> f (Maybe a, Maybe b)
{-# INLINE [1] cpadZip #-}
cpadZip :: f a -> f b -> f (Maybe a, Maybe b)
cpadZip = (These a b -> (Maybe a, Maybe b))
-> f a -> f b -> f (Maybe a, Maybe b)
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith (Maybe a
-> Maybe b -> These (Maybe a) (Maybe b) -> (Maybe a, Maybe b)
forall a b. a -> b -> These a b -> (a, b)
fromThese Maybe a
forall a. Maybe a
Nothing Maybe b
forall a. Maybe a
Nothing (These (Maybe a) (Maybe b) -> (Maybe a, Maybe b))
-> (These a b -> These (Maybe a) (Maybe b))
-> These a b
-> (Maybe a, Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe a)
-> (b -> Maybe b) -> These a b -> These (Maybe a) (Maybe b)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> Maybe a
forall a. a -> Maybe a
Just b -> Maybe b
forall a. a -> Maybe a
Just)

cpadZipWith
  :: (CSemialign f, Dom f a, Dom f b, Dom f c)
  => (Maybe a -> Maybe b -> c)
  -> f a -> f b -> f c
{-# INLINE [1] cpadZipWith #-}
cpadZipWith :: (Maybe a -> Maybe b -> c) -> f a -> f b -> f c
cpadZipWith Maybe a -> Maybe b -> c
f = (These a b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
(CSemialign f, Dom f a, Dom f b, Dom f c) =>
(These a b -> c) -> f a -> f b -> f c
calignWith ((These a b -> c) -> f a -> f b -> f c)
-> (These a b -> c) -> f a -> f b -> f c
forall a b. (a -> b) -> a -> b
$
  (Maybe a -> Maybe b -> c) -> (Maybe a, Maybe b) -> c
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Maybe a -> Maybe b -> c
f ((Maybe a, Maybe b) -> c)
-> (These a b -> (Maybe a, Maybe b)) -> These a b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe a
-> Maybe b -> These (Maybe a) (Maybe b) -> (Maybe a, Maybe b)
forall a b. a -> b -> These a b -> (a, b)
fromThese Maybe a
forall a. Maybe a
Nothing Maybe b
forall a. Maybe a
Nothing (These (Maybe a) (Maybe b) -> (Maybe a, Maybe b))
-> (These a b -> These (Maybe a) (Maybe b))
-> These a b
-> (Maybe a, Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe a)
-> (b -> Maybe b) -> These a b -> These (Maybe a) (Maybe b)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> Maybe a
forall a. a -> Maybe a
Just b -> Maybe b
forall a. a -> Maybe a
Just