Make List.zipWith as lazy as expected

Author: aspiwack

src/Data/List/Linear.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE LinearTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE NoImplicitPrelude #-}
@@ -353,18 +354,34 @@ zip3 :: (Consumable a, Consumable b, Consumable c) => [a] %1 -> [b] %1 -> [c] %1
 zip3 = zipWith3 (,,)
 
 zipWith :: (Consumable a, Consumable b) => (a %1 -> b %1 -> c) -> [a] %1 -> [b] %1 -> [c]
-zipWith f xs ys =
-  zipWith' f xs ys & \(ret, leftovers) ->
-    leftovers `lseq` ret
+zipWith f =
+  zipWithk f (:) [] consume2 consume2
+  where
+    consume2 :: forall x y z. (Consumable x, Consumable y) => x %1 -> y %1 -> [z]
+    consume2 x y = x `lseq` y `lseq` []
 
 -- | Same as 'zipWith', but returns the leftovers instead of consuming them.
+-- Because the leftovers are returned at toplevel, @zipWith'@ is pretty strict:
+-- forcing the first cons cell of the returned list forces all the recursive
+-- calls.
 zipWith' :: (a %1 -> b %1 -> c) -> [a] %1 -> [b] %1 -> ([c], Maybe (Either (NonEmpty a) (NonEmpty b)))
-zipWith' _ [] [] = ([], Nothing)
-zipWith' _ (a : as) [] = ([], Just (Left (a :| as)))
-zipWith' _ [] (b : bs) = ([], Just (Right (b :| bs)))
-zipWith' f (a : as) (b : bs) =
-  case zipWith' f as bs of
-    (cs, rest) -> (f a b : cs, rest)
+zipWith' f =
+  zipWithk
+    f
+    (\c !(cs, rest) -> ((c : cs), rest))
+    ([], Nothing)
+    (\a as -> ([], Just (Left (a :| as))))
+    (\b bs -> ([], Just (Right (b :| bs))))
+
+zipWithk :: forall r a b c. (a %1 -> b %1 -> c) -> (c %1 -> r %1 -> r) -> r -> (a %1 -> [a] %1 -> r) -> (b %1 -> [b] %1 -> r) -> [a] %1 -> [b] %1 -> r
+zipWithk f cons nil lefta leftb =
+  go
+  where
+    go :: [a] %1 -> [b] %1 -> r
+    go [] [] = nil
+    go (a : as) [] = lefta a as
+    go [] (b : bs) = leftb b bs
+    go (a : as) (b : bs) = cons (f a b) (go as bs)
 
 zipWith3 :: forall a b c d. (Consumable a, Consumable b, Consumable c) => (a %1 -> b %1 -> c %1 -> d) -> [a] %1 -> [b] %1 -> [c] %1 -> [d]
 zipWith3 _ [] ys zs = (ys, zs) `lseq` []