{-# 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.Word ()
import Data.List.NonEmpty (NonEmpty (..))
import Data.Word (Word8)
import Numeric.Natural (Natural)

instance Memoizable Natural where
  memoize :: (Natural -> b) -> Natural -> b
memoize Natural -> b
f = (NonEmpty Word8 -> b) -> NonEmpty Word8 -> b
forall a b. Memoizable a => (a -> b) -> a -> b
memoize (Natural -> b
f (Natural -> b)
-> (NonEmpty Word8 -> Natural) -> NonEmpty Word8 -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Word8 -> Natural
wordsToNat) (NonEmpty Word8 -> b)
-> (Natural -> NonEmpty Word8) -> Natural -> b
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 = [Word8] -> Natural -> NonEmpty Word8
forall a a. (Bits a, Integral a, Num a) => [a] -> a -> NonEmpty a
go []
  where
    go :: [a] -> a -> NonEmpty a
go [a]
acc a
n =
      if a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0xff
        then a
w a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
acc
        else [a] -> a -> NonEmpty a
go (a
w a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
8 a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
      where
        w :: a
w = a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n

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