{- |
This module is intended for high-performance programming. Concretely, the 'ExpandCase' class can be
used to force GHC to avoid compiling t'Control.Monad.Reader.Reader', t'Data.Monoid.Endo' or t'Control.Monad.State.State'-based code to unnecessary closures.
You can observe its usage in "Mikan.TypeChecking.Free.Generic". For a smaller example and some
explanation, consider the following.

@
f :: Bool -> Endo Int
f b = case b of
  True  -> mempty
  False -> Endo (+ 10)
@

The desired @-O1@ output should be the following (ignoring newtype casts):

@
f = \\b n -> case b of
  True  -> n
  False -> n + 10
@

A typical undesired output would be

@
f = \\b -> case b of
  True  -> \\n -> n
  False -> \\n -> n + 10
@

Returning closures can be better or worse, depending on the program context. However, in
high-performance situations we almost never want to return closures, and GHC is not nearly reliable
enough at getting rid of the closures.

Using 'ExpandCase', we can write code as follows.

@
f :: Bool -> Endo Int
f b = expand \\ret -> case b of
  True  -> ret mempty
  False -> ret $ Endo (+10)
@

Here, 'expand' immediately introduces a lambda abstraction, and the @ret@ continuation applies the
"body" of the definition to the freshly abstracted variable. Hence, we get something like the
following as an intermediate piece of Core:

@
f :: Bool -> Endo Int
f b = Endo \\n -> case b of
  True  -> appEndo mempty n
  False -> appEndo (Endo (+10)) n
@

which is then reliably optimized to

@
f :: Bool -> Endo Int
f b = Endo \\n -> case b of
  True  -> n
  False -> n + 10
@

NOTE: if you want to use this module, it is very strongly recommended that you check the Core of
your code!
-}

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
module Mikan.Utils.ExpandCase
  ( ExpandCase(..)
    -- * Deriving helpers
  , DontExpand(..)
    -- * Re-exports
  , LiftedRep
  )
  where

import Data.Monoid
import GHC.Exts (oneShot, RuntimeRep(..), TYPE, LiftedRep)
import Data.Strict.Tuple

class ExpandCase rep a | a -> rep where
  type Result rep a :: TYPE rep
  expand :: ((a -> Result rep a) -> Result rep a) -> a

--------------------------------------------------------------------------------
-- Deriving helpers

newtype DontExpand a = DontExpand { forall a. DontExpand a -> a
unDontExpand :: a }

instance ExpandCase LiftedRep (DontExpand a) where
  type Result LiftedRep (DontExpand a) = a

  {-# INLINE expand #-}
  expand :: ((DontExpand a -> Result LiftedRep (DontExpand a))
 -> Result LiftedRep (DontExpand a))
-> DontExpand a
expand (DontExpand a -> Result LiftedRep (DontExpand a))
-> Result LiftedRep (DontExpand a)
k = a -> DontExpand a
forall a. a -> DontExpand a
DontExpand ((DontExpand a -> Result LiftedRep (DontExpand a))
-> Result LiftedRep (DontExpand a)
k DontExpand a -> a
DontExpand a -> Result LiftedRep (DontExpand a)
forall a. DontExpand a -> a
unDontExpand)

--------------------------------------------------------------------------------
-- Instances

deriving via DontExpand Any        instance ExpandCase LiftedRep Any
deriving via DontExpand All        instance ExpandCase LiftedRep All
deriving via DontExpand Bool       instance ExpandCase LiftedRep Bool
deriving via DontExpand Int        instance ExpandCase LiftedRep Int
deriving via DontExpand ()         instance ExpandCase LiftedRep ()
deriving via DontExpand (Pair a b) instance ExpandCase LiftedRep (Pair a b)
deriving via DontExpand (IO a)     instance ExpandCase LiftedRep (IO a)
deriving via DontExpand [a]        instance ExpandCase LiftedRep [a]
deriving via DontExpand (Maybe a)  instance ExpandCase LiftedRep (Maybe a)

instance ExpandCase LiftedRep (Endo a) where
  type Result LiftedRep (Endo a) = a
  {-# INLINE expand #-}
  expand :: ((Endo a -> Result LiftedRep (Endo a))
 -> Result LiftedRep (Endo a))
-> Endo a
expand (Endo a -> Result LiftedRep (Endo a)) -> Result LiftedRep (Endo a)
k = (a -> a) -> Endo a
forall a. (a -> a) -> Endo a
Endo ((a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
oneShot \a
a -> (Endo a -> Result LiftedRep (Endo a)) -> Result LiftedRep (Endo a)
k ((Endo a -> a) -> Endo a -> a
forall a b. (a -> b) -> a -> b
oneShot \Endo a
act -> Endo a -> a -> a
forall a. Endo a -> a -> a
appEndo Endo a
act a
a))