From 9de2abbf1d723246e2c0ac41cb7a929a683a36be Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Mon, 9 May 2022 16:21:15 -0400
Subject: [PATCH] [#265] Make SuperComposition less brittle

* Use overlapping instances instead of incoherent ones. Fixes #265.

* Make the first argument of `...` a function unconditionally, before
  instance selection. This can theoretically improve inference slightly,
  though it probably doesn't have much impact in practice.
---
 CHANGES.md              |  4 +++
 benchmark/Main.hs       |  4 +--
 src/Universum/VarArg.hs | 66 +++++++++++++++++++++++++++++++++--------
 3 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index ece3451..3b4148c 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,6 +1,10 @@
 Unreleased
 =====
 
+* [#265](https://github.com/serokell/universum/issues/265):
+  Make `SuperComposition` inference less brittle, and give it four
+  type parameters.
+
 * [#252](https://github.com/serokell/universum/pull/252):
   Remove `Option` re-export. Use `Maybe` instead.
 
diff --git a/benchmark/Main.hs b/benchmark/Main.hs
index b507209..c82abc7 100644
--- a/benchmark/Main.hs
+++ b/benchmark/Main.hs
@@ -112,9 +112,9 @@ bgroupSuperComposition = bgroup "(...)"
  where
   super10 :: [()] -> Bool
   super10 = null
-        ... ((: []) ... Unsafe.head ... pure ... Unsafe.head
+        ... (: []) ... Unsafe.head ... (pure :: () -> [()]) ... Unsafe.head
         ... (: [(), (), (), ()]) ... Unsafe.head ... (: []) ... Unsafe.head
-        ... (: [()]) ... Unsafe.head ... (: [(), ()]) ... Unsafe.head :: [()] -> [()])
+        ... (: [()]) ... Unsafe.head ... (: [(), ()]) ... Unsafe.head
 
   norm10 = null
          . (: []) . Unsafe.head . pure . Unsafe.head
diff --git a/src/Universum/VarArg.hs b/src/Universum/VarArg.hs
index 18f7cae..88df350 100644
--- a/src/Universum/VarArg.hs
+++ b/src/Universum/VarArg.hs
@@ -1,9 +1,14 @@
+{-# LANGUAGE AllowAmbiguousTypes    #-}
+{-# LANGUAGE DataKinds              #-}
 {-# LANGUAGE FlexibleContexts       #-}
 {-# LANGUAGE FlexibleInstances      #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE MultiParamTypeClasses  #-}
 {-# LANGUAGE Safe                   #-}
+{-# LANGUAGE ScopedTypeVariables    #-}
+{-# LANGUAGE TypeApplications       #-}
 {-# LANGUAGE TypeFamilies           #-}
+{-# LANGUAGE TypeOperators          #-}
 {-# LANGUAGE UndecidableInstances   #-}
 
 -- | Provides operator of variable-arguments function composition.
@@ -11,6 +16,9 @@
 module Universum.VarArg
        ( SuperComposition(..)
        ) where
+import Data.Type.Bool (Not, type (||))
+import Data.Type.Equality (type (==))
+import Prelude (Bool (..))
 
 -- $setup
 -- >>> import Universum.Base ((+))
@@ -19,17 +27,25 @@ module Universum.VarArg
 -- >>> import Data.List (zip5)
 
 -- | This type class allows to implement variadic composition operator.
-class SuperComposition a b c | a b -> c where
-    -- | Allows to apply function to result of another function with multiple
+class SuperComposition x y b r | x y b -> r where
+    -- | Applies a function to the result of another function with multiple
     -- arguments.
     --
-    -- >>> (show ... (+)) 1 2
+    -- >>> (show ... (+)) (1 :: Int) 2
     -- "3"
-    -- >>> show ... 5
+    -- >>> (show ... (+)) 1 2 :: String
+    -- "3"
+    -- >>> show ... (5 :: Int)
+    -- "5"
+    -- >>> show ... 5 :: String
     -- "5"
     -- >>> (null ... zip5) [1] [2] [3] [] [5]
     -- True
     --
+    -- Note that the type checker needs to have enough information on hand to deduce
+    -- the appropriate arity for the second argument, which explains the need for explicit
+    -- types in some examples above.
+    --
     -- Inspired by <http://stackoverflow.com/questions/9656797/variadic-compose-function>.
     --
     -- ==== Performance
@@ -45,16 +61,42 @@ class SuperComposition a b c | a b -> c where
     -- disappear due to very general inferred type. However, functions without type
     -- specification but with applied @INLINE@ pragma are fast again.
     --
-    (...) :: a -> b -> c
+    (...) :: (x -> y) -> b -> r
 
 infixl 8 ...
 
-instance {-# INCOHERENT #-} (a ~ c, r ~ b) =>
-         SuperComposition (a -> b) c r where
-    f ... g = f g
-    {-# INLINE (...) #-}
+-- The implementation is a bit tricky to get right. See #265 for how things can go wrong.
+-- The basic idea is that we can commit to using the base case if we know we've reached
+-- a result of the right type *or* we know that we don't have any arrows left. Similarly,
+-- we can commit to using the recursive case if we know we don't yet have a result of the
+-- right type *or* we know that we have more arrows we can use.
+
+type family IsArrow b where
+  IsArrow (_ -> _) = 'True
+  IsArrow _ = 'False
+
+-- | Can we use the base case?
+type PlainApplication y b r = y == r || Not (IsArrow b)
+
+-- | Can we use the recursive case?
+type Composing y b r = Not (y == r) || IsArrow b
+
+class SuperComposition' (plainApplication :: Bool) (composing :: Bool) x y b r | x y b -> r where
+    comp :: (x -> y) -> b -> r
+
+instance (x ~ b, y ~ r) =>
+         SuperComposition' 'True composing x y b r where
+    comp f = f
+    {-# INLINE comp #-}
+
+instance {-# INCOHERENT #-} (b ~ (b1 -> b'), r ~ (b1 -> r'), SuperComposition x y b' r') =>
+         SuperComposition' plainApplication 'True x y b r where
+    (f `comp` g) c = f ... g c
+    {-# INLINE comp #-}
 
-instance {-# INCOHERENT #-} (SuperComposition (a -> b) d r1, r ~ (c -> r1)) =>
-         SuperComposition (a -> b) (c -> d) r where
-    (f ... g) c = f ... g c
+instance ( pa ~ PlainApplication y b r
+         , co ~ Composing y b r
+         , SuperComposition' pa co x y b r) =>
+         SuperComposition x y b r where
+    (...) = comp @pa @co
     {-# INLINE (...) #-}