{-# OPTIONS_GHC -Wno-orphans #-}

module Data.Function.FastMemo.Natural () where

import Data.Bits (shiftL, shiftR, (.|.))
import Data.Foldable (foldl')
import Data.Function.FastMemo.Class (Memoizable (..))
import Data.Function.FastMemo.List ()
import Data.Function.FastMemo.Word ()
import Data.List.NonEmpty (NonEmpty (..))
import Data.Word (Word8)
import Numeric.Natural (Natural)

instance Memoizable Natural where
  memoize :: forall b. (Natural -> b) -> Natural -> b
memoize Natural -> b
f = forall a b. Memoizable a => (a -> b) -> a -> b
memoize (Natural -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Word8 -> Natural
wordsToNat) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> NonEmpty Word8
natToWords

-- A slightly weird encoding that assigns unique values to each of [0], [0,0], [0,0,0]...
natToWords :: Natural -> NonEmpty Word8
natToWords :: Natural -> NonEmpty Word8
natToWords = forall {t} {a}.
(Bits t, Integral t, Num a) =>
[a] -> t -> NonEmpty a
go []
  where
    go :: [a] -> t -> NonEmpty a
go [a]
acc t
n =
      if t
n forall a. Ord a => a -> a -> Bool
<= t
0xff
        then a
w forall a. a -> [a] -> NonEmpty a
:| [a]
acc
        else [a] -> t -> NonEmpty a
go (a
w forall a. a -> [a] -> [a]
: [a]
acc) (t
n forall a. Bits a => a -> Int -> a
`shiftR` Int
8 forall a. Num a => a -> a -> a
- t
1)
      where
        w :: a
w = forall a b. (Integral a, Num b) => a -> b
fromIntegral t
n

wordsToNat :: NonEmpty Word8 -> Natural
wordsToNat :: NonEmpty Word8 -> Natural
wordsToNat (Word8
w0 :| [Word8]
ws0) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Natural
acc Word8
w -> (Natural
acc forall a. Num a => a -> a -> a
+ Natural
1) forall a. Bits a => a -> Int -> a
`shiftL` Int
8 forall a. Bits a => a -> a -> a
.|. forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w0) [Word8]
ws0