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

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

import Data.List.NonEmpty (NonEmpty)
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 = (Rep a Any -> b) -> Rep a Any -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (a -> b
f (a -> b) -> (Rep a Any -> a) -> Rep a Any -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rep a Any -> a
forall a x. Generic a => Rep a x -> a
to) (Rep a Any -> b) -> (a -> Rep a Any) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Rep a Any
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 :: (M1 i c f p -> b) -> M1 i c f p -> b
gMemoize M1 i c f p -> b
f = (f p -> b) -> f p -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (M1 i c f p -> b
f (M1 i c f p -> b) -> (f p -> M1 i c f p) -> f p -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f p -> M1 i c f p
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1) (f p -> b) -> (M1 i c f p -> f p) -> M1 i c f p -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. M1 i c f p -> f p
forall i (c :: Meta) k (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1

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

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

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

instance (GMemoize a, GMemoize b) => GMemoize (a :*: b) where
  gMemoize :: ((:*:) a b p -> b) -> (:*:) a b p -> b
gMemoize (:*:) a b p -> b
f =
    let f' :: a p -> b p -> b
f' = (a p -> b p -> b) -> a p -> b p -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (\a p
x -> (b p -> b) -> b p -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize (\b p
y -> (:*:) a b p -> b
f (a p
x a p -> b p -> (:*:) a b p
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 :: ((:+:) a b p -> b) -> (:+:) a b p -> b
gMemoize (:+:) a b p -> b
f =
    let fL :: a p -> b
fL = (a p -> b) -> a p -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize ((:+:) a b p -> b
f ((:+:) a b p -> b) -> (a p -> (:+:) a b p) -> a p -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a p -> (:+:) a b p
forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1)
        fR :: b p -> b
fR = (b p -> b) -> b p -> b
forall (a :: * -> *) p b. GMemoize a => (a p -> b) -> a p -> b
gMemoize ((:+:) a b p -> b
f ((:+:) a b p -> b) -> (b p -> (:+:) a b p) -> b p -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b p -> (:+:) a b p
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)

instance Memoizable a => Memoizable [a]

instance Memoizable a => Memoizable (NonEmpty a)