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 :: forall n. Integral n => 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 {t}. Integral t => t -> t -> [t]
dr n
base n
i
    where
      dr :: t -> t -> [t]
dr t
_ t
0 = []
      dr t
b t
x = case n
base of
                n
1 -> t -> [t] -> [t]
forall i a. Integral i => i -> [a] -> [a]
genericTake t
x ([t] -> [t]) -> [t] -> [t]
forall a b. (a -> b) -> a -> b
$ t -> [t]
forall a. a -> [a]
repeat t
1
                n
_ -> let (t
rest, t
lastDigit) = t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
quotRem t
x t
b
                     in t
lastDigit t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> t -> [t]
dr t
b t
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 :: forall n. Integral n => 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 :: forall {t}. Integral t => t -> t -> [t]
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 :: forall {t}. Integral t => t -> t -> [t]
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 {t}. Integral t => t -> t -> [t]
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 :: forall n. Integral n => n -> [n] -> n
unDigits n
base = (n -> n -> n) -> n -> [n] -> n
forall b a. (b -> a -> b) -> b -> [a] -> b
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 {t}. Integral t => t -> t -> [t]
digits Integer
b) Integer
i