{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeOperators #-}

module Data.Function.FastMemo.Class (Memoizable (..)) where

import GHC.Generics

class Memoizable a where
  memoize :: (a -> b) -> a -> b
  default memoize :: (Generic a, GMemoize (Rep a)) => (a -> b) -> a -> b
  memoize a -> b
f = forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a x. Generic a => Rep a x -> a
to) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a x. Generic a => a -> Rep a x
from

class GMemoize a where
  gMemoize :: (a p -> b) -> a p -> b

instance GMemoize f => GMemoize (M1 i c f) where
  gMemoize :: forall p b. (M1 i c f p -> b) -> M1 i c f p -> b
gMemoize M1 i c f p -> b
f = forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (M1 i c f p -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1

instance GMemoize V1 where
  gMemoize :: forall p b. (V1 p -> b) -> V1 p -> b
gMemoize V1 p -> b
_f = \case {}

instance GMemoize U1 where
  gMemoize :: forall p b. (U1 p -> b) -> U1 p -> b
gMemoize U1 p -> b
f = let fu :: b
fu = U1 p -> b
f forall k (p :: k). U1 p
U1 in \U1 p
U1 -> b
fu

instance Memoizable c => GMemoize (K1 i c) where
  gMemoize :: forall p b. (K1 i c p -> b) -> K1 i c p -> b
gMemoize K1 i c p -> b
f = forall a b. Memoizable a => (a -> b) -> a -> b
memoize (K1 i c p -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i c (p :: k). c -> K1 i c p
K1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i c (p :: k). K1 i c p -> c
unK1

instance (GMemoize a, GMemoize b) => GMemoize (a :*: b) where
  gMemoize :: forall p b. ((:*:) a b p -> b) -> (:*:) a b p -> b
gMemoize (:*:) a b p -> b
f =
    let f' :: a p -> b p -> b
f' = forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (\a p
x -> forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (\b p
y -> (:*:) a b p -> b
f (a p
x forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: b p
y)))
     in \(a p
x :*: b p
y) -> a p -> b p -> b
f' a p
x b p
y

instance (GMemoize a, GMemoize b) => GMemoize (a :+: b) where
  gMemoize :: forall p b. ((:+:) a b p -> b) -> (:+:) a b p -> b
gMemoize (:+:) a b p -> b
f =
    let fL :: a p -> b
fL = forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize ((:+:) a b p -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1)
        fR :: b p -> b
fR = forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize ((:+:) a b p -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1)
     in \case
          L1 a p
x -> a p -> b
fL a p
x
          R1 b p
x -> b p -> b
fR b p
x

instance (Memoizable a, Memoizable b) => Memoizable (a, b)