{-# 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((.&.), xor, shiftR, shiftL), FiniteBits(finiteBitSize) )
import Data.Int ( Int8, Int16, Int32, Int64 )
import Data.Word ( Word8, Word16, Word32, Word64 )
import Numeric.Natural ( Natural )

-- $setup
-- >>> :set -XNegativeLiterals -XScopedTypeVariables -XFlexibleContexts
-- >>> import Data.Word
-- >>> import Data.Int
-- >>> import Numeric.Natural
-- >>> import Test.QuickCheck.Instances.Natural
{-|
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 = signed -> unsigned
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    ((signed
s signed -> Int -> signed
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) signed -> signed -> signed
forall a. Bits a => a -> a -> a
`xor` (signed
s signed -> Int -> signed
forall a. Bits a => a -> Int -> a
`shiftR` (signed -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize signed
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))

  {-# INLINE zigZag #-}
  zagZig :: unsigned -> signed
  default zagZig :: (Bits unsigned) => unsigned -> signed
  zagZig unsigned
u = unsigned -> signed
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((unsigned
u unsigned -> Int -> unsigned
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) unsigned -> unsigned -> unsigned
forall a. Bits a => a -> a -> a
`xor` unsigned -> unsigned
forall a. Num a => a -> a
negate (unsigned
u unsigned -> unsigned -> unsigned
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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
x Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Int
1
    | Bool
otherwise = Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer -> Integer
forall a. Num a => a -> a
negate (Integer
x Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1

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