{-# LANGUAGE LambdaCase #-}

module Data.Function.FastMemo.Util (memoizeFixedLen) where

import Data.Function.FastMemo.Class (Memoizable, memoize)
import GHC.Stack (HasCallStack)

-- | Memoize a function on a list whose length is predetermined.
--
-- If called on a larger list, it will truncate; on a smaller list, it will throw an error
memoizeFixedLen :: (HasCallStack, Memoizable a) => Int -> ([a] -> b) -> [a] -> b
memoizeFixedLen :: Int -> ([a] -> b) -> [a] -> b
memoizeFixedLen Int
n [a] -> b
f
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = b -> [a] -> b
forall a b. a -> b -> a
const ([a] -> b
f [])
  | Bool
otherwise =
    let f' :: a -> [a] -> b
f' = (a -> [a] -> b) -> a -> [a] -> b
forall a b. Memoizable a => (a -> b) -> a -> b
memoize ((a -> [a] -> b) -> a -> [a] -> b)
-> (a -> [a] -> b) -> a -> [a] -> b
forall a b. (a -> b) -> a -> b
$ \a
x -> Int -> ([a] -> b) -> [a] -> b
forall a b.
(HasCallStack, Memoizable a) =>
Int -> ([a] -> b) -> [a] -> b
memoizeFixedLen (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([a] -> b
f ([a] -> b) -> ([a] -> [a]) -> [a] -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:))
     in \case
          [] -> [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"List too short"
          a
x : [a]
xs -> a -> [a] -> b
f' a
x [a]
xs