{-# LANGUAGE ExplicitForAll      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax       #-}
{-# LANGUAGE TypeApplications    #-}

-- | Immutable lazy tables of functions over bounded enumerations.

-- Function calls are stored as thunks and not evaluated until accessed.

--

-- The underlying representation is an 'Data.Array.Array' rather than a search

-- tree. This provides /O(1)/ lookup but means that the range of keys should not

-- be very large, as in the case of an 'Int'-like type.

module Data.Enum.Memo
  ( -- * Memoization

    memoize
    -- * Higher-order

  , memoize2
  , memoize3
  , memoize4
  , memoize5
  ) where

import Prelude hiding (lookup)

import Data.Array ((!), array)

-- | Memoize a function with a single argument.

memoize :: ∀ k v. (Bounded k, Enum k) => (k -> v) -> k -> v
memoize :: (k -> v) -> k -> v
memoize k -> v
f = case (Int, Int) -> [(Int, v)] -> Array Int v
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int, Int)
bounds [(Int, v)]
vals of
    Array Int v
memo -> \k
k -> Array Int v
memo Array Int v -> Int -> v
forall i e. Ix i => Array i e -> i -> e
! k -> Int
forall a. Enum a => a -> Int
fromEnum k
k
  where
    bounds :: (Int, Int)
bounds = (k -> Int
forall a. Enum a => a -> Int
fromEnum @k k
forall a. Bounded a => a
minBound, k -> Int
forall a. Enum a => a -> Int
fromEnum @k k
forall a. Bounded a => a
maxBound)
    vals :: [(Int, v)]
vals   = [(k -> Int
forall a. Enum a => a -> Int
fromEnum k
k, k -> v
f k
k) | k
k <- [k
forall a. Bounded a => a
minBound..k
forall a. Bounded a => a
maxBound]]

-- | Memoize a function with two arguments.

memoize2 :: ∀ k1 k2 v.
            ( Bounded k1, Enum k1
            , Bounded k2, Enum k2
            )
         => (k1 -> k2 -> v) -> k1 -> k2 -> v
memoize2 :: (k1 -> k2 -> v) -> k1 -> k2 -> v
memoize2 k1 -> k2 -> v
f = case ((Int, Int), (Int, Int)) -> [((Int, Int), v)] -> Array (Int, Int) v
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int, Int), (Int, Int))
bounds [((Int, Int), v)]
vals of
    Array (Int, Int) v
memo -> \k1
k1 k2
k2 -> Array (Int, Int) v
memo Array (Int, Int) v -> (Int, Int) -> v
forall i e. Ix i => Array i e -> i -> e
!
            ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
            , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
            )
  where
    bounds :: ((Int, Int), (Int, Int))
bounds = ( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
minBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
minBound
               )
             , ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
maxBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
maxBound
               )
             )
    vals :: [((Int, Int), v)]
vals   = [( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
                , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
                )
              , k1 -> k2 -> v
f k1
k1 k2
k2
              ) | k1
k1 <- [k1
forall a. Bounded a => a
minBound..k1
forall a. Bounded a => a
maxBound]
                , k2
k2 <- [k2
forall a. Bounded a => a
minBound..k2
forall a. Bounded a => a
maxBound]]

-- | Memoize a function with three arguments.

memoize3 :: ∀ k1 k2 k3 v.
            ( Bounded k1, Enum k1
            , Bounded k2, Enum k2
            , Bounded k3, Enum k3
            )
         => (k1 -> k2 -> k3 -> v) -> k1 -> k2 -> k3 -> v
memoize3 :: (k1 -> k2 -> k3 -> v) -> k1 -> k2 -> k3 -> v
memoize3 k1 -> k2 -> k3 -> v
f = case ((Int, Int, Int), (Int, Int, Int))
-> [((Int, Int, Int), v)] -> Array (Int, Int, Int) v
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int, Int, Int), (Int, Int, Int))
bounds [((Int, Int, Int), v)]
vals of
    Array (Int, Int, Int) v
memo -> \k1
k1 k2
k2 k3
k3 -> Array (Int, Int, Int) v
memo Array (Int, Int, Int) v -> (Int, Int, Int) -> v
forall i e. Ix i => Array i e -> i -> e
!
            ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
            , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
            , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
            )
  where
    bounds :: ((Int, Int, Int), (Int, Int, Int))
bounds = ( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
minBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
minBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
minBound
               )
             , ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
maxBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
maxBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
maxBound
               )
             )
    vals :: [((Int, Int, Int), v)]
vals   = [( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
                , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
                , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
                )
              , k1 -> k2 -> k3 -> v
f k1
k1 k2
k2 k3
k3
              ) | k1
k1 <- [k1
forall a. Bounded a => a
minBound..k1
forall a. Bounded a => a
maxBound]
                , k2
k2 <- [k2
forall a. Bounded a => a
minBound..k2
forall a. Bounded a => a
maxBound]
                , k3
k3 <- [k3
forall a. Bounded a => a
minBound..k3
forall a. Bounded a => a
maxBound]]

-- | Memoize a function with four arguments.

memoize4 :: ∀ k1 k2 k3 k4 v.
            ( Bounded k1, Enum k1
            , Bounded k2, Enum k2
            , Bounded k3, Enum k3
            , Bounded k4, Enum k4
            )
         => (k1 -> k2 -> k3 -> k4 -> v) -> k1 -> k2 -> k3 -> k4 -> v
memoize4 :: (k1 -> k2 -> k3 -> k4 -> v) -> k1 -> k2 -> k3 -> k4 -> v
memoize4 k1 -> k2 -> k3 -> k4 -> v
f = case ((Int, Int, Int, Int), (Int, Int, Int, Int))
-> [((Int, Int, Int, Int), v)] -> Array (Int, Int, Int, Int) v
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int, Int, Int, Int), (Int, Int, Int, Int))
bounds [((Int, Int, Int, Int), v)]
vals of
    Array (Int, Int, Int, Int) v
memo -> \k1
k1 k2
k2 k3
k3 k4
k4 -> Array (Int, Int, Int, Int) v
memo Array (Int, Int, Int, Int) v -> (Int, Int, Int, Int) -> v
forall i e. Ix i => Array i e -> i -> e
!
            ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
            , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
            , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
            , k4 -> Int
forall a. Enum a => a -> Int
fromEnum k4
k4
            )
  where
    bounds :: ((Int, Int, Int, Int), (Int, Int, Int, Int))
bounds = ( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
minBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
minBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
minBound
               , k4 -> Int
forall a. Enum a => a -> Int
fromEnum @k4 k4
forall a. Bounded a => a
minBound
               )
             , ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
maxBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
maxBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
maxBound
               , k4 -> Int
forall a. Enum a => a -> Int
fromEnum @k4 k4
forall a. Bounded a => a
maxBound
               )
             )
    vals :: [((Int, Int, Int, Int), v)]
vals   = [( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
                , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
                , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
                , k4 -> Int
forall a. Enum a => a -> Int
fromEnum k4
k4
                )
              , k1 -> k2 -> k3 -> k4 -> v
f k1
k1 k2
k2 k3
k3 k4
k4
              ) | k1
k1 <- [k1
forall a. Bounded a => a
minBound..k1
forall a. Bounded a => a
maxBound]
                , k2
k2 <- [k2
forall a. Bounded a => a
minBound..k2
forall a. Bounded a => a
maxBound]
                , k3
k3 <- [k3
forall a. Bounded a => a
minBound..k3
forall a. Bounded a => a
maxBound]
                , k4
k4 <- [k4
forall a. Bounded a => a
minBound..k4
forall a. Bounded a => a
maxBound]]

-- | Memoize a function with five arguments.

memoize5 :: ∀ k1 k2 k3 k4 k5 v.
            ( Bounded k1, Enum k1
            , Bounded k2, Enum k2
            , Bounded k3, Enum k3
            , Bounded k4, Enum k4
            , Bounded k5, Enum k5
            )
         => (k1 -> k2 -> k3 -> k4 -> k5 -> v) -> k1 -> k2 -> k3 -> k4 -> k5 -> v
memoize5 :: (k1 -> k2 -> k3 -> k4 -> k5 -> v)
-> k1 -> k2 -> k3 -> k4 -> k5 -> v
memoize5 k1 -> k2 -> k3 -> k4 -> k5 -> v
f = case ((Int, Int, Int, Int, Int), (Int, Int, Int, Int, Int))
-> [((Int, Int, Int, Int, Int), v)]
-> Array (Int, Int, Int, Int, Int) v
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int, Int, Int, Int, Int), (Int, Int, Int, Int, Int))
bounds [((Int, Int, Int, Int, Int), v)]
vals of
    Array (Int, Int, Int, Int, Int) v
memo -> \k1
k1 k2
k2 k3
k3 k4
k4 k5
k5 -> Array (Int, Int, Int, Int, Int) v
memo Array (Int, Int, Int, Int, Int) v -> (Int, Int, Int, Int, Int) -> v
forall i e. Ix i => Array i e -> i -> e
!
            ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
            , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
            , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
            , k4 -> Int
forall a. Enum a => a -> Int
fromEnum k4
k4
            , k5 -> Int
forall a. Enum a => a -> Int
fromEnum k5
k5
            )
  where
    bounds :: ((Int, Int, Int, Int, Int), (Int, Int, Int, Int, Int))
bounds = ( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
minBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
minBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
minBound
               , k4 -> Int
forall a. Enum a => a -> Int
fromEnum @k4 k4
forall a. Bounded a => a
minBound
               , k5 -> Int
forall a. Enum a => a -> Int
fromEnum @k5 k5
forall a. Bounded a => a
minBound
               )
             , ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum @k1 k1
forall a. Bounded a => a
maxBound
               , k2 -> Int
forall a. Enum a => a -> Int
fromEnum @k2 k2
forall a. Bounded a => a
maxBound
               , k3 -> Int
forall a. Enum a => a -> Int
fromEnum @k3 k3
forall a. Bounded a => a
maxBound
               , k4 -> Int
forall a. Enum a => a -> Int
fromEnum @k4 k4
forall a. Bounded a => a
maxBound
               , k5 -> Int
forall a. Enum a => a -> Int
fromEnum @k5 k5
forall a. Bounded a => a
maxBound
               )
             )
    vals :: [((Int, Int, Int, Int, Int), v)]
vals   = [( ( k1 -> Int
forall a. Enum a => a -> Int
fromEnum k1
k1
                , k2 -> Int
forall a. Enum a => a -> Int
fromEnum k2
k2
                , k3 -> Int
forall a. Enum a => a -> Int
fromEnum k3
k3
                , k4 -> Int
forall a. Enum a => a -> Int
fromEnum k4
k4
                , k5 -> Int
forall a. Enum a => a -> Int
fromEnum k5
k5
                )
              , k1 -> k2 -> k3 -> k4 -> k5 -> v
f k1
k1 k2
k2 k3
k3 k4
k4 k5
k5
              ) | k1
k1 <- [k1
forall a. Bounded a => a
minBound..k1
forall a. Bounded a => a
maxBound]
                , k2
k2 <- [k2
forall a. Bounded a => a
minBound..k2
forall a. Bounded a => a
maxBound]
                , k3
k3 <- [k3
forall a. Bounded a => a
minBound..k3
forall a. Bounded a => a
maxBound]
                , k4
k4 <- [k4
forall a. Bounded a => a
minBound..k4
forall a. Bounded a => a
maxBound]
                , k5
k5 <- [k5
forall a. Bounded a => a
minBound..k5
forall a. Bounded a => a
maxBound]]