Unzip helpers

semialign/semialign.cabal

@@ -61,12 +61,17 @@ library
     Data.Zip
 
   other-modules:    Data.Semialign.Internal
+                  , Data.Semialign.Internal.Tuples
 
   -- ghc boot libs
   build-depends:
       base          >=4.5.1.0 && <4.16
     , containers    >=0.4.2.1 && <0.7
     , transformers  >=0.3.0.0 && <0.7
+  if impl (ghc < 9.0.1)
+    build-depends:
+      -- For noinline
+      ghc-prim
 
   -- These
   build-depends:    these >=1.1.1.1 && <1.2
@@ -79,10 +84,7 @@ library
     , tagged                         >=0.8.6    && <0.9
     , unordered-containers           >=0.2.8.0  && <0.3
     , vector                         >=0.12.0.2 && <0.13
-
-  -- base shims
-  if !impl(ghc >=8.2)
-    build-depends: bifunctors >=5.5.4 && <5.6
+    , bifunctors                     >=5.5.4 && <5.6
 
   if !impl(ghc >=8.0)
     build-depends:

semialign/src/Data/Semialign/Internal.hs

@@ -1,4 +1,7 @@
 {-# LANGUAGE CPP                        #-}
+{-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE DeriveFoldable             #-}
+{-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE FunctionalDependencies     #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -13,12 +16,13 @@ module Data.Semialign.Internal where
 import Prelude
        (Bool (..), Either (..), Eq (..), Functor (fmap), Int, Maybe (..),
        Monad (..), Ord (..), Ordering (..), String, error, flip, fst, id,
-       maybe, snd, uncurry, ($), (++), (.))
+       maybe, snd, uncurry, ($), (++), (.), Traversable, Foldable)
 
 import qualified Prelude as Prelude
 
 import Control.Applicative               (ZipList (..), pure, (<$>))
 import Data.Bifunctor                    (Bifunctor (..))
+import Data.Biapplicative                (traverseBia)
 import Data.Functor.Compose              (Compose (..))
 import Data.Functor.Identity             (Identity (..))
 import Data.Functor.Product              (Product (..))
@@ -74,6 +78,8 @@ import           Data.IntMap (IntMap)
 import qualified Data.IntMap as IntMap
 #endif
 
+import Data.Semialign.Internal.Tuples (SBPair (..), LBPair (..))
+
 import Data.These
 import Data.These.Combinators
 
@@ -577,6 +583,58 @@ instance (Ord k) => Align (Map k) where
 instance Ord k => Unalign (Map k) where
     unalign xs = (Map.mapMaybe justHere xs, Map.mapMaybe justThere xs)
 
+newtype UnzipStrictSpineStrictPairs t a =
+  UnzipStrictSpineStrictPairs { getUnzipStrictSpineStrictPairs :: t a }
+  deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
+
+instance (Zip t, Traversable t) => Unzip (UnzipStrictSpineStrictPairs t) where
+  unzipWith = unzipWithStrictSpineStrictPairs
+
+newtype UnzipStrictSpineLazyPairs t a =
+  UnzipStrictSpineLazyPairs { getUnzipStrictSpineLazyPairs :: t a }
+  deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
+
+instance (Zip t, Traversable t) => Unzip (UnzipStrictSpineLazyPairs t) where
+  unzipWith = unzipWithStrictSpineLazyPairs
+  unzip = unzipStrictSpineLazyPairs
+
+newtype UnzipLazySpineLazyPairs t a =
+  UnzipLazySpineLazyPairs { getUnzipLazySpineLazyPairs :: t a }
+  deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
+
+instance (Zip t, Traversable t) => Unzip (UnzipLazySpineLazyPairs t) where
+  unzipWith = unzipWithLazySpineLazyPairs
+
+unzipWithStrictSpineStrictPairs :: Traversable t
+  => (c -> (a, b)) -> t c -> (t a, t b)
+unzipWithStrictSpineStrictPairs f = unSBPair . traverseBia (SBPair . f)
+
+unzipWithStrictSpineLazyPairs :: Traversable t
+  => (c -> (a, b)) -> t c -> (t a, t b)
+unzipWithStrictSpineLazyPairs f = unSBPair . traverseBia (SBPair . foo)
+  where
+    foo c = let
+      {-# NOINLINE fc #-}
+      {-# NOINLINE a #-}
+      {-# NOINLINE b #-}
+      fc = f c
+      (a, b) = fc
+      in (a, b)
+
+unzipStrictSpineLazyPairs :: Traversable t
+  => t (a, b) -> (t a, t b)
+unzipStrictSpineLazyPairs = unSBPair . traverseBia (SBPair . foo)
+  where
+    foo ab = let
+      {-# NOINLINE a #-}
+      {-# NOINLINE b #-}
+      (a, b) = ab
+      in (a, b)
+
+unzipWithLazySpineLazyPairs :: Traversable t
+  => (c -> (a, b)) -> t c -> (t a, t b)
+unzipWithLazySpineLazyPairs f = unLBPair . traverseBia (LBPair . f)
+
 instance Ord k => Unzip (Map k) where unzip = unzipDefault
 
 instance Ord k => Zip (Map k) where

semialign/src/Data/Semialign/Internal/Tuples.hs

@@ -0,0 +1,132 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE Trustworthy #-}
+module Data.Semialign.Internal.Tuples
+  ( SBPair (..)
+  , LBPair (..)
+  , Solo (..)
+  , getSolo
+  ) where
+
+import Data.Bifunctor (Bifunctor (..))
+import Data.Biapplicative (Biapplicative (..))
+
+#if !MIN_VERSION_base(4,8,0)
+import Control.Applicative (Applicative (..))
+#endif
+
+#if MIN_VERSION_base(4,15,0)
+import GHC.Exts (noinline)
+#elif MIN_VERSION_ghc_prim(0,5,1)
+import GHC.Magic (noinline)
+#endif
+
+-- A copy of (,) with a stricter bimap.
+newtype SBPair a b = SBPair { unSBPair :: (a, b) }
+
+instance Bifunctor SBPair where
+  bimap f g (SBPair (a, b)) = SBPair (f a, g b)
+
+instance Biapplicative SBPair where
+  bipure a b = SBPair (a, b)
+  biliftA2 f g (SBPair (a, b)) (SBPair (c, d)) =
+    SBPair (f a c, g b d)
+
+-- A copy of (,) with a lazier biliftA2
+newtype LBPair a b = LBPair { unLBPair :: (a, b) }
+
+instance Bifunctor LBPair where
+  bimap = bimapLB
+
+bimapLB :: (a -> c) -> (b -> d) -> LBPair a b -> LBPair c d
+bimapLB f g (LBPair ab) = LBPair (f a, g b)
+    where
+      -- This stuff can be really touchy, so we're extra careful.
+      -- We want a and b to be actual selector thunks. If their
+      -- definitions inline, then they won't be. Why do we say
+      -- noinline ab? That may be a bit belt-and-suspenders, but
+      -- I've been bitten in the past. The concern is that GHC
+      -- could see
+      --
+      --   bimapLB f g p@(LBPair (e1, e2))
+      --
+      -- and decide to do something like
+      --
+      --   let (a, _) = p
+      --   in LBPair (f a, g e2)
+      --
+      -- I don't remember the details, but something similar happened
+      -- when defining Data.List.transpose, so I'll just be careful
+      -- until it's proven unnecessary.
+      {-# NOINLINE a #-}
+      {-# NOINLINE b #-}
+      (a, b) = noinline ab
+{-# NOINLINE [1] bimapLB #-}
+
+-- Optimize when we can, being sure to expand both sides.
+-- Hopefully these rules can't break the selector thunks.
+{-# RULES
+"bimap/known" forall f g a b. bimapLB f g (LBPair (a, b)) = LBPair (f a, g b)
+ #-}
+
+instance Biapplicative LBPair where
+  bipure a b = LBPair (a, b)
+  biliftA2 = biliftA2LB
+
+biliftA2LB :: (a -> c -> e) -> (b -> d -> f) -> LBPair a b -> LBPair c d -> LBPair e f
+biliftA2LB f g (LBPair ab) (LBPair cd) = LBPair (f a c, g b d)
+    where
+      {-# NOINLINE a #-}
+      {-# NOINLINE b #-}
+      {-# NOINLINE c #-}
+      {-# NOINLINE d #-}
+      (a, b) = noinline ab
+      (c, d) = noinline cd
+{-# NOINLINE [1] biliftA2LB #-}
+
+biliftA2LBkl :: (a -> c -> e) -> (b -> d -> f) -> a -> b -> LBPair c d -> LBPair e f
+biliftA2LBkl f g a b (LBPair cd) = LBPair (f a c, g b d)
+    where
+      {-# NOINLINE c #-}
+      {-# NOINLINE d #-}
+      (c, d) = noinline cd
+{-# NOINLINE [1] biliftA2LBkl #-}
+
+biliftA2LBkr :: (a -> c -> e) -> (b -> d -> f) -> LBPair a b -> c -> d -> LBPair e f
+biliftA2LBkr f g (LBPair ab) c d = LBPair (f a c, g b d)
+    where
+      {-# NOINLINE a #-}
+      {-# NOINLINE b #-}
+      (a, b) = noinline ab
+{-# NOINLINE [1] biliftA2LBkr #-}
+
+{-# RULES
+"biliftA2/knownl" forall f g a b cd. biliftA2LB f g (LBPair (a, b)) cd
+  = biliftA2LBkl f g a b cd
+"biliftA2/knownlr" forall f g a b c d. biliftA2LBkl f g a b (LBPair (c, d))
+  = LBPair (f a c, g b d)
+"biliftA2/knownr" forall f g ab c d. biliftA2LB f g ab (LBPair (c, d))
+  = biliftA2LBkr f g ab c d
+"biliftA2/knownrl" forall f g a b c d. biliftA2LBkr f g (LBPair (a, b)) c d
+  = LBPair (f a c, g b d)
+ #-}
+
+-- ----------
+-- Compat stuff.
+
+-- As of GHC 9.0, Solo is not exported from base (it's stuck in ghc-prim).
+-- Hopefully this will be sorted by 9.2, and it will definitely be sorted by
+-- 9.4. I'd rather avoid an unconditional dependency on ghc-prim, especially
+-- when we just need two instances and one of them is derived.
+data Solo a = Solo { getSolo :: a }
+  deriving Functor
+
+instance Applicative Solo where
+  pure = Solo
+  Solo f <*> Solo a = Solo (f a)
+
+#if !MIN_VERSION_ghc_prim(0,5,1)
+{-# NOINLINE noinline #-}
+noinline :: a -> a
+noinline a = a
+#endif

semialign/src/Data/Zip.hs

@@ -10,6 +10,14 @@ module Data.Zip (
     Unzip (..),
     unzipDefault,
     Zippy (..),
+    -- * Unzip definition helpers
+    UnzipStrictSpineStrictPairs (..),
+    UnzipStrictSpineLazyPairs (..),
+    UnzipLazySpineLazyPairs (..),
+    unzipWithStrictSpineStrictPairs,
+    unzipWithStrictSpineLazyPairs,
+    unzipStrictSpineLazyPairs,
+    unzipWithLazySpineLazyPairs,
     ) where
 
 import Control.Applicative (Applicative (..))