module Data.Digits (mDigits, digits, mDigitsRev, digitsRev, unDigits, prop_digitsRoundTrip) where

import Test.QuickCheck
import Data.Maybe (fromJust)
import Data.List (genericTake)

-- | Returns the digits of a positive integer as a Maybe list, in reverse order
--   or Nothing if a zero or negative base is given
--   This is slightly more efficient than in forward order.
mDigitsRev :: Integral n
    => n         -- ^ The base to use.
    -> n         -- ^ The number to convert to digit form.
    -> Maybe [n] -- ^ Nothing or Just the digits of the number in list form, in reverse.
mDigitsRev :: n -> n -> Maybe [n]
mDigitsRev n
base n
i = if n
base n -> n -> Bool
forall a. Ord a => a -> a -> Bool
< n
1
                    then Maybe [n]
forall a. Maybe a
Nothing -- We do not support zero or negative bases
                    else [n] -> Maybe [n]
forall a. a -> Maybe a
Just ([n] -> Maybe [n]) -> [n] -> Maybe [n]
forall a b. (a -> b) -> a -> b
$ n -> n -> [n]
forall a. Integral a => a -> a -> [a]
dr n
base n
i
    where
      dr :: a -> a -> [a]
dr a
_ a
0 = []
      dr a
b a
x = case n
base of
                n
1 -> a -> [a] -> [a]
forall i a. Integral i => i -> [a] -> [a]
genericTake a
x ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ a -> [a]
forall a. a -> [a]
repeat a
1
                n
_ -> let (a
rest, a
lastDigit) = a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem a
x a
b
                     in a
lastDigit a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> a -> [a]
dr a
b a
rest

-- | Returns the digits of a positive integer as a Maybe list.
--   or Nothing if a zero or negative base is given
mDigits :: Integral n
    => n -- ^ The base to use.
    -> n -- ^ The number to convert to digit form.
    -> Maybe [n] -- ^ Nothing or Just the digits of the number in list form
mDigits :: n -> n -> Maybe [n]
mDigits n
base n
i = [n] -> [n]
forall a. [a] -> [a]
reverse ([n] -> [n]) -> Maybe [n] -> Maybe [n]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> n -> n -> Maybe [n]
forall n. Integral n => n -> n -> Maybe [n]
mDigitsRev n
base n
i

-- | Returns the digits of a positive integer as a list, in reverse order.
--   Throws an error if given a zero or negative base.
digitsRev :: Integral n
    => n   -- ^ The base to use.
    -> n   -- ^ The number to convert to digit from.
    -> [n] -- ^ The digits of the number in list from, in reverse.
digitsRev :: n -> n -> [n]
digitsRev n
base = Maybe [n] -> [n]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [n] -> [n]) -> (n -> Maybe [n]) -> n -> [n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> n -> Maybe [n]
forall n. Integral n => n -> n -> Maybe [n]
mDigitsRev n
base

-- | Returns the digits of a positive integer as a list.
--   Throws an error if given a zero or negative base.
digits :: Integral n
    => n   -- ^ The base to use (typically 10).
    -> n   -- ^ The number to convert to digit form.
    -> [n] -- ^ Either Nothing or the digits of the number in list form.
digits :: n -> n -> [n]
digits n
base = [n] -> [n]
forall a. [a] -> [a]
reverse ([n] -> [n]) -> (n -> [n]) -> n -> [n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> n -> [n]
forall a. Integral a => a -> a -> [a]
digitsRev n
base

-- | Takes a list of digits, and converts them back into a positive integer.
unDigits :: Integral n
    => n   -- ^ The base to use.
    -> [n] -- ^ The digits of the number in list form.
    -> n   -- ^ The original number.
unDigits :: n -> [n] -> n
unDigits n
base = (n -> n -> n) -> n -> [n] -> n
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ n
a n
b -> n
a n -> n -> n
forall a. Num a => a -> a -> a
* n
base n -> n -> n
forall a. Num a => a -> a -> a
+ n
b) n
0

-- | unDigits . digits should be the identity, in any positive base.
prop_digitsRoundTrip
    :: Integer -- ^ The integer to test.
    -> Integer -- ^ The base to use.
    -> Property
prop_digitsRoundTrip :: Integer -> Integer -> Property
prop_digitsRoundTrip Integer
i Integer
b = Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 Bool -> Property -> Property
forall prop. Testable prop => Bool -> prop -> Property
==> Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 Bool -> Bool -> Property
forall prop. Testable prop => Bool -> prop -> Property
==> Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== (Integer -> [Integer] -> Integer
forall n. Integral n => n -> [n] -> n
unDigits Integer
b ([Integer] -> Integer)
-> (Integer -> [Integer]) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer -> [Integer]
forall a. Integral a => a -> a -> [a]
digits Integer
b) Integer
i