Skip to content

Commit 0ace0d9

Browse files
committed
Unzip definitions
1 parent 7a9277d commit 0ace0d9

File tree

4 files changed

+202
-12
lines changed

4 files changed

+202
-12
lines changed

semialign/semialign.cabal

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,12 +61,17 @@ library
6161
Data.Zip
6262

6363
other-modules: Data.Semialign.Internal
64+
, Data.Semialign.Internal.Tuples
6465

6566
-- ghc boot libs
6667
build-depends:
6768
base >=4.5.1.0 && <4.16
6869
, containers >=0.4.2.1 && <0.7
6970
, transformers >=0.3.0.0 && <0.7
71+
if impl (ghc < 9.0.1)
72+
build-depends:
73+
-- For noinline
74+
ghc-prim
7075

7176
-- These
7277
build-depends: these >=1.1.1.1 && <1.2

semialign/src/Data/Semialign/Internal.hs

Lines changed: 57 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
{-# LANGUAGE CPP #-}
22
{-# LANGUAGE DeriveFunctor #-}
3+
{-# LANGUAGE DeriveFoldable #-}
4+
{-# LANGUAGE DeriveTraversable #-}
35
{-# LANGUAGE FlexibleInstances #-}
46
{-# LANGUAGE FunctionalDependencies #-}
57
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -14,13 +16,13 @@ module Data.Semialign.Internal where
1416
import Prelude
1517
(Bool (..), Either (..), Eq (..), Functor (fmap), Int, Maybe (..),
1618
Monad (..), Ord (..), Ordering (..), String, error, flip, fst, id,
17-
maybe, snd, uncurry, ($), (++), (.))
19+
maybe, snd, uncurry, ($), (++), (.), Traversable, Foldable)
1820

1921
import qualified Prelude as Prelude
2022

2123
import Control.Applicative (ZipList (..), pure, (<$>))
2224
import Data.Bifunctor (Bifunctor (..))
23-
import Data.Biapplicative (Biapplicative (..), traverseBia)
25+
import Data.Biapplicative (traverseBia)
2426
import Data.Functor.Compose (Compose (..))
2527
import Data.Functor.Identity (Identity (..))
2628
import Data.Functor.Product (Product (..))
@@ -76,6 +78,8 @@ import Data.IntMap (IntMap)
7678
import qualified Data.IntMap as IntMap
7779
#endif
7880

81+
import Data.Semialign.Internal.Tuples (SBPair (..), LBPair (..))
82+
7983
import Data.These
8084
import Data.These.Combinators
8185

@@ -579,16 +583,57 @@ instance (Ord k) => Align (Map k) where
579583
instance Ord k => Unalign (Map k) where
580584
unalign xs = (Map.mapMaybe justHere xs, Map.mapMaybe justThere xs)
581585

582-
-- A copy of (,) with a stricter bimap.
583-
newtype SBPair a b = SBPair { unSBPair :: (a, b) }
584-
585-
instance Bifunctor SBPair where
586-
bimap f g (SBPair (a, b)) = SBPair (f a, g b)
587-
588-
instance Biapplicative SBPair where
589-
bipure a b = SBPair (a, b)
590-
biliftA2 f g (SBPair (a, b)) (SBPair (c, d)) =
591-
SBPair (f a c, g b d)
586+
newtype UnzipStrictSpineStrictPairs t a =
587+
UnzipStrictSpineStrictPairs { getUnzipStrictSpineStrictPairs :: t a }
588+
deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
589+
590+
instance (Zip t, Traversable t) => Unzip (UnzipStrictSpineStrictPairs t) where
591+
unzipWith = unzipWithStrictSpineStrictPairs
592+
593+
newtype UnzipStrictSpineLazyPairs t a =
594+
UnzipStrictSpineLazyPairs { getUnzipStrictSpineLazyPairs :: t a }
595+
deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
596+
597+
instance (Zip t, Traversable t) => Unzip (UnzipStrictSpineLazyPairs t) where
598+
unzipWith = unzipWithStrictSpineLazyPairs
599+
unzip = unzipStrictSpineLazyPairs
600+
601+
newtype UnzipLazySpineLazyPairs t a =
602+
UnzipLazySpineLazyPairs { getUnzipLazySpineLazyPairs :: t a }
603+
deriving (Functor, Foldable, Traversable, Semialign, Align, Zip)
604+
605+
instance (Zip t, Traversable t) => Unzip (UnzipLazySpineLazyPairs t) where
606+
unzipWith = unzipWithLazySpineLazyPairs
607+
608+
unzipWithStrictSpineStrictPairs :: Traversable t
609+
=> (c -> (a, b)) -> t c -> (t a, t b)
610+
unzipWithStrictSpineStrictPairs f = unSBPair . traverseBia (SBPair . f)
611+
612+
unzipWithStrictSpineLazyPairs :: Traversable t
613+
=> (c -> (a, b)) -> t c -> (t a, t b)
614+
unzipWithStrictSpineLazyPairs f = unSBPair . traverseBia (SBPair . foo)
615+
where
616+
foo c = let
617+
{-# NOINLINE fc #-}
618+
{-# NOINLINE a #-}
619+
{-# NOINLINE b #-}
620+
fc = f c
621+
(a, b) = fc
622+
in (a, b)
623+
624+
unzipStrictSpineLazyPairs :: Traversable t
625+
=> t (a, b) -> (t a, t b)
626+
unzipStrictSpineLazyPairs = unSBPair . traverseBia (SBPair . foo)
627+
where
628+
foo ab = let
629+
{-# NOINLINE a #-}
630+
{-# NOINLINE b #-}
631+
(a, b) = ab
632+
in (a, b)
633+
634+
unzipWithLazySpineLazyPairs :: Traversable t
635+
=> (c -> (a, b)) -> t c -> (t a, t b)
636+
unzipWithLazySpineLazyPairs f = unLBPair . traverseBia (LBPair . f)
592637

593638
instance Ord k => Unzip (Map k) where unzip = unzipDefault
594639

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
{-# LANGUAGE CPP #-}
2+
{-# LANGUAGE DeriveFunctor #-}
3+
{-# LANGUAGE Trustworthy #-}
4+
module Data.Semialign.Internal.Tuples
5+
( SBPair (..)
6+
, LBPair (..)
7+
, Solo (..)
8+
, getSolo
9+
) where
10+
11+
import Data.Bifunctor (Bifunctor (..))
12+
import Data.Biapplicative (Biapplicative (..))
13+
14+
#if !MIN_VERSION_base(4,8,0)
15+
import Control.Applicative (Applicative (..))
16+
#endif
17+
18+
#if MIN_VERSION_base(4,15,0)
19+
import GHC.Exts (noinline)
20+
#elif MIN_VERSION_ghc_prim(0,5,1)
21+
import GHC.Magic (noinline)
22+
#endif
23+
24+
-- A copy of (,) with a stricter bimap.
25+
newtype SBPair a b = SBPair { unSBPair :: (a, b) }
26+
27+
instance Bifunctor SBPair where
28+
bimap f g (SBPair (a, b)) = SBPair (f a, g b)
29+
30+
instance Biapplicative SBPair where
31+
bipure a b = SBPair (a, b)
32+
biliftA2 f g (SBPair (a, b)) (SBPair (c, d)) =
33+
SBPair (f a c, g b d)
34+
35+
-- A copy of (,) with a lazier biliftA2
36+
newtype LBPair a b = LBPair { unLBPair :: (a, b) }
37+
38+
instance Bifunctor LBPair where
39+
bimap = bimapLB
40+
41+
bimapLB :: (a -> c) -> (b -> d) -> LBPair a b -> LBPair c d
42+
bimapLB f g (LBPair ab) = LBPair (f a, g b)
43+
where
44+
-- This stuff can be really touchy, so we're extra careful.
45+
-- We want a and b to be actual selector thunks. If their
46+
-- definitions inline, then they won't be. Why do we say
47+
-- noinline ab? That may be a bit belt-and-suspenders, but
48+
-- I've been bitten in the past. The concern is that GHC
49+
-- could see
50+
--
51+
-- bimapLB f g p@(LBPair (e1, e2))
52+
--
53+
-- and decide to do something like
54+
--
55+
-- let (a, _) = p
56+
-- in LBPair (f a, g e2)
57+
--
58+
-- I don't remember the details, but something similar happened
59+
-- when defining Data.List.transpose, so I'll just be careful
60+
-- until it's proven unnecessary.
61+
{-# NOINLINE a #-}
62+
{-# NOINLINE b #-}
63+
(a, b) = noinline ab
64+
{-# NOINLINE [1] bimapLB #-}
65+
66+
-- Optimize when we can, being sure to expand both sides.
67+
-- Hopefully these rules can't break the selector thunks.
68+
{-# RULES
69+
"bimap/known" forall f g a b. bimapLB f g (LBPair (a, b)) = LBPair (f a, g b)
70+
#-}
71+
72+
instance Biapplicative LBPair where
73+
bipure a b = LBPair (a, b)
74+
biliftA2 = biliftA2LB
75+
76+
biliftA2LB :: (a -> c -> e) -> (b -> d -> f) -> LBPair a b -> LBPair c d -> LBPair e f
77+
biliftA2LB f g (LBPair ab) (LBPair cd) = LBPair (f a c, g b d)
78+
where
79+
{-# NOINLINE a #-}
80+
{-# NOINLINE b #-}
81+
{-# NOINLINE c #-}
82+
{-# NOINLINE d #-}
83+
(a, b) = noinline ab
84+
(c, d) = noinline cd
85+
{-# NOINLINE [1] biliftA2LB #-}
86+
87+
biliftA2LBkl :: (a -> c -> e) -> (b -> d -> f) -> a -> b -> LBPair c d -> LBPair e f
88+
biliftA2LBkl f g a b (LBPair cd) = LBPair (f a c, g b d)
89+
where
90+
{-# NOINLINE c #-}
91+
{-# NOINLINE d #-}
92+
(c, d) = noinline cd
93+
{-# NOINLINE [1] biliftA2LBkl #-}
94+
95+
biliftA2LBkr :: (a -> c -> e) -> (b -> d -> f) -> LBPair a b -> c -> d -> LBPair e f
96+
biliftA2LBkr f g (LBPair ab) c d = LBPair (f a c, g b d)
97+
where
98+
{-# NOINLINE a #-}
99+
{-# NOINLINE b #-}
100+
(a, b) = noinline ab
101+
{-# NOINLINE [1] biliftA2LBkr #-}
102+
103+
{-# RULES
104+
"biliftA2/knownl" forall f g a b cd. biliftA2LB f g (LBPair (a, b)) cd
105+
= biliftA2LBkl f g a b cd
106+
"biliftA2/knownlr" forall f g a b c d. biliftA2LBkl f g a b (LBPair (c, d))
107+
= LBPair (f a c, g b d)
108+
"biliftA2/knownr" forall f g ab c d. biliftA2LB f g ab (LBPair (c, d))
109+
= biliftA2LBkr f g ab c d
110+
"biliftA2/knownrl" forall f g a b c d. biliftA2LBkr f g (LBPair (a, b)) c d
111+
= LBPair (f a c, g b d)
112+
#-}
113+
114+
-- ----------
115+
-- Compat stuff.
116+
117+
-- As of GHC 9.0, Solo is not exported from base (it's stuck in ghc-prim).
118+
-- Hopefully this will be sorted by 9.2, and it will definitely be sorted by
119+
-- 9.4. I'd rather avoid an unconditional dependency on ghc-prim, especially
120+
-- when we just need two instances and one of them is derived.
121+
data Solo a = Solo { getSolo :: a }
122+
deriving Functor
123+
124+
instance Applicative Solo where
125+
pure = Solo
126+
Solo f <*> Solo a = Solo (f a)
127+
128+
#if !MIN_VERSION_ghc_prim(0,5,1)
129+
{-# NOINLINE noinline #-}
130+
noinline :: a -> a
131+
noinline a = a
132+
#endif

semialign/src/Data/Zip.hs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,14 @@ module Data.Zip (
1010
Unzip (..),
1111
unzipDefault,
1212
Zippy (..),
13+
-- * Unzip definition helpers
14+
UnzipStrictSpineStrictPairs (..),
15+
UnzipStrictSpineLazyPairs (..),
16+
UnzipLazySpineLazyPairs (..),
17+
unzipWithStrictSpineStrictPairs,
18+
unzipWithStrictSpineLazyPairs,
19+
unzipStrictSpineLazyPairs,
20+
unzipWithLazySpineLazyPairs,
1321
) where
1422

1523
import Control.Applicative (Applicative (..))

0 commit comments

Comments
 (0)