From 05f4e88ae45511c262a074ea0e344709d577f7cb Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Sun, 22 Aug 2021 15:48:16 -0400
Subject: [PATCH] Unzip Map and IntMap more strictly

The previous `Map` (same for `IntMap` throughout) instance
would first map eagerly over tha `Map`, producing an entire
`Map` full of thunks to apply the `unzipWith` function. Then
it would build two more entire `Map`s full of thunks to select
components of each pair. We can do it eagerly using
`Data.Biapplicative.traverseBia` and a simple pair bifunctor.
---
 semialign/semialign.cabal                |  5 +----
 semialign/src/Data/Semialign/Internal.hs | 26 ++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/semialign/semialign.cabal b/semialign/semialign.cabal
index 8a06793..3c0f5df 100644
--- a/semialign/semialign.cabal
+++ b/semialign/semialign.cabal
@@ -79,10 +79,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:
diff --git a/semialign/src/Data/Semialign/Internal.hs b/semialign/src/Data/Semialign/Internal.hs
index d19d223..3b6d7e3 100644
--- a/semialign/src/Data/Semialign/Internal.hs
+++ b/semialign/src/Data/Semialign/Internal.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP                        #-}
+{-# LANGUAGE DeriveFunctor              #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE FunctionalDependencies     #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -19,6 +20,7 @@ import qualified Prelude as Prelude
 
 import Control.Applicative               (ZipList (..), pure, (<$>))
 import Data.Bifunctor                    (Bifunctor (..))
+import Data.Biapplicative                (Biapplicative (..), traverseBia)
 import Data.Functor.Compose              (Compose (..))
 import Data.Functor.Identity             (Identity (..))
 import Data.Functor.Product              (Product (..))
@@ -577,7 +579,25 @@ instance (Ord k) => Align (Map k) where
 instance Ord k => Unalign (Map k) where
     unalign xs = (Map.mapMaybe justHere xs, Map.mapMaybe justThere xs)
 
-instance Ord k => Unzip (Map k) where unzip = unzipDefault
+-- 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)
+
+instance Ord k => Unzip (Map k) where
+  -- Map has a strict spine, so we have to build a whole one at
+  -- once. The default instance would first build an entire
+  -- Map filled with thunks, each of which will produce a pair,
+  -- and then build two maps, each filled with thunks to extract
+  -- a value from the pair. Sounds like a mess. Let's just do
+  -- the whole calculation eagerly.
+  unzipWith f = unSBPair . traverseBia (SBPair . f)
 
 instance Ord k => Zip (Map k) where
     zipWith = Map.intersectionWith
@@ -601,7 +621,9 @@ instance Align IntMap where
 instance Unalign IntMap where
     unalign xs = (IntMap.mapMaybe justHere xs, IntMap.mapMaybe justThere xs)
 
-instance Unzip IntMap where unzip = unzipDefault
+instance Unzip IntMap where
+  -- See notes at the Map instance
+  unzipWith f = unSBPair . traverseBia (SBPair . f)
 
 instance Zip IntMap where
     zipWith = IntMap.intersectionWith