{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE ScopedTypeVariables    #-}

-- |<https://gist.github.com/mfuerstenau/ba870a29e16536fdbaba ZigZag encoding> of signed integrals.
module Data.ZigZag (ZigZag(..)) where

import           Data.Bits       (Bits (shiftL, shiftR, xor, (.&.)),
                                  FiniteBits (finiteBitSize))
import           Data.Int        (Int16, Int32, Int64, Int8)
import           Data.Word       (Word16, Word32, Word64, Word8)
import           Numeric.Natural (Natural)

-- $setup
-- >>> :set -XNegativeLiterals -XScopedTypeVariables -XFlexibleContexts
-- >>> import Data.Word
-- >>> import Data.Int
-- >>> import Numeric.Natural
-- >>> import Test.QuickCheck.Arbitrary
-- >>> instance Arbitrary Natural where arbitrary = arbitrarySizedNatural; shrink    = shrinkIntegral

{-|
Convert between a signed integral and the corresponding ZigZag encoded unsigned integral (e.g. between Int8 and Word8 or Integral and Natural).

Allow conversion only between compatible types, invalid conversions produce a type error:

@
zigZag (-1::Int64) :: Word32
...
... Couldn't match type ...
...
@
>>> zigZag (0::Int8)
0

>>> zigZag (-1::Int16)
1

>>> zigZag (1::Int32)
2

>>> zigZag (-2::Int16)
3

>>> zigZag (-50::Integer)
99

>>> zigZag (50::Integer)
100

>>> zigZag (64::Integer)
128

>>> zigZag (-256::Integer)
511

>>> zigZag (256::Integer)
512

>>> map zigZag [-3..3::Integer]
[5,3,1,0,2,4,6]

>>> map zagZig [0..6::Word8]
[0,-1,1,-2,2,-3,3]

prop> \(f::Integer) -> zagZig (zigZag f) == f
+++ OK, passed 100 tests.

prop> \(f::Natural) -> zigZag (zagZig f) == f
+++ OK, passed 100 tests.

prop> \(f::Int8) -> zagZig (zigZag f) == f
+++ OK, passed 100 tests.

prop> \(f::Word8) -> zigZag (zagZig f) == f
+++ OK, passed 100 tests.

prop> \(s::Int8) -> zigZag s == fromIntegral (zigZag (fromIntegral s :: Integer))
+++ OK, passed 100 tests.

prop> \(u::Word8) -> zagZig u == fromIntegral (zagZig (fromIntegral u :: Natural))
+++ OK, passed 100 tests.

prop> \(f::Int64) -> zagZig (zigZag f) == f
+++ OK, passed 100 tests.

prop> \(f::Word64) -> zigZag (zagZig f) == f
+++ OK, passed 100 tests.

prop> \(s::Int64) -> zigZag s == fromIntegral (zigZag (fromIntegral s :: Integer))
+++ OK, passed 100 tests.

prop> \(u::Word64) -> zagZig u == fromIntegral (zagZig (fromIntegral u :: Natural))
+++ OK, passed 100 tests.
-}
class (Integral signed, Integral unsigned)
  => ZigZag signed unsigned | unsigned -> signed, signed -> unsigned where
  zigZag :: signed -> unsigned
  default zigZag :: FiniteBits signed => signed -> unsigned
  zigZag signed
s = forall a b. (Integral a, Num b) => a -> b
fromIntegral
    ((signed
s forall a. Bits a => a -> Int -> a
`shiftL` Int
1) forall a. Bits a => a -> a -> a
`xor` (signed
s forall a. Bits a => a -> Int -> a
`shiftR` (forall b. FiniteBits b => b -> Int
finiteBitSize signed
s forall a. Num a => a -> a -> a
- Int
1)))

  {-# INLINE zigZag #-}
  zagZig :: unsigned -> signed
  default zagZig :: (Bits unsigned) => unsigned -> signed
  zagZig unsigned
u = forall a b. (Integral a, Num b) => a -> b
fromIntegral ((unsigned
u forall a. Bits a => a -> Int -> a
`shiftR` Int
1) forall a. Bits a => a -> a -> a
`xor` forall a. Num a => a -> a
negate (unsigned
u forall a. Bits a => a -> a -> a
.&. unsigned
1))

  -- default zagZig :: (Bits signed) => unsigned -> signed
  -- zagZig u = let (s::signed) = fromIntegral u in ((s `shiftR` 1) `xor` (negate (s .&. 1)))
  {-# INLINE zagZig #-}

instance ZigZag Int8 Word8

instance ZigZag Int16 Word16

instance ZigZag Int32 Word32

instance ZigZag Int64 Word64

instance ZigZag Integer Natural where
  zigZag :: Integer -> Natural
zigZag Integer
x
    | Integer
x forall a. Ord a => a -> a -> Bool
>= Integer
0 = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Integer
x forall a. Bits a => a -> Int -> a
`shiftL` Int
1
    | Bool
otherwise = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
negate (Integer
x forall a. Bits a => a -> Int -> a
`shiftL` Int
1) forall a. Num a => a -> a -> a
- Integer
1

  zagZig :: Natural -> Integer
zagZig Natural
u = let s :: Integer
s = forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
u
             in ((Integer
s forall a. Bits a => a -> Int -> a
`shiftR` Int
1) forall a. Bits a => a -> a -> a
`xor` forall a. Num a => a -> a
negate (Integer
s forall a. Bits a => a -> a -> a
.&. Integer
1))