{-# LANGUAGE TypeApplications #-}

{- | Zigzag encoding maps signed integers to unsigned integers so that numbers
with a small absolute value (for instance, -1) have a small varint encoded
value too. It does this in a way that "zig-zags" back and forth through the
positive and negative integers, so that -1 is encoded as 1, 1 is encoded as
2, -2 is encoded as 3, and so on.

> zigzag(n) = { 2 * n       if 0 <= n
>             { -2 * n - 1  if n < 0

This description was adapted from
https://developers.google.com/protocol-buffers/docs/encoding#signed-ints
which is released under https://creativecommons.org/licenses/by/4.0/
-}
module Data.Word.Zigzag
  ( toZigzag
  , fromZigzag
  , toZigzagNative
  , fromZigzagNative
  , toZigzag32
  , fromZigzag32
  , toZigzag64
  , fromZigzag64
  ) where

import Data.Bits (complement, unsafeShiftL, unsafeShiftR, xor, (.&.))
import Data.Int (Int32, Int64)
import Data.Word (Word32, Word64)
import Numeric.Natural (Natural)

{- | Encode a big integer with zigzag.

If you know the size of the data, it is likely more efficient to use one of
'toZigzagNative', 'toZigzag32', or 'toZigzag64'.
-}
toZigzag :: Integer -> Natural
toZigzag :: Integer -> Natural
toZigzag Integer
n
  | Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
n = Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
n
  | 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
2) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1

{- | Decode a zigzag-encoded big ingeter.

If you know the size of the data, it is likely more efficient to use one of
'fromZigzagNative', 'fromZigzag32', or 'fromZigzag64'.
-}
fromZigzag :: Natural -> Integer
fromZigzag :: Natural -> Integer
fromZigzag Natural
n
  | Natural -> Bool
forall a. Integral a => a -> Bool
even Natural
n = Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Natural
n Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`div` Natural
2
  | Bool
otherwise = Integer -> Integer
forall a. Num a => a -> a
negate (Integer -> Integer) -> (Natural -> Integer) -> Natural -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1) Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`div` Natural
2

{- | Encode a native-size integer with zigzag.

In C, this is:

> (n << 1) ^ (n >> (CHAR_BIT * sizeof(int) - 1))
-}
toZigzagNative :: Int -> Word
toZigzagNative :: Int -> Word
toZigzagNative = Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Word) -> (Int -> Word64) -> Int -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
toZigzag64 (Int64 -> Word64) -> (Int -> Int64) -> Int -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{- | Decode a native-size zigzag-encoded integer.

In C, this is:

> (n >> 1) ^ (~(n & 1) + 1)
-}
fromZigzagNative :: Word -> Int
fromZigzagNative :: Word -> Int
fromZigzagNative = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> (Word -> Int64) -> Word -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Int64
fromZigzag64 (Word64 -> Int64) -> (Word -> Word64) -> Word -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{- | Encode a 32-bit integer with zigzag.

In C, this is:

> (n << 1) ^ (n >> 31)
-}
toZigzag32 :: Int32 -> Word32
toZigzag32 :: Int32 -> Word32
toZigzag32 = Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Word32) -> (Int32 -> Word64) -> Int32 -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
toZigzag64 (Int64 -> Word64) -> (Int32 -> Int64) -> Int32 -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{- | Decode a 32-bit zigzag-encoded integer.

In C, this is:

> (n >> 1) ^ (~(n & 1) + 1)
-}
fromZigzag32 :: Word32 -> Int32
fromZigzag32 :: Word32 -> Int32
fromZigzag32 = Int64 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int32) -> (Word32 -> Int64) -> Word32 -> Int32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Int64
fromZigzag64 (Word64 -> Int64) -> (Word32 -> Word64) -> Word32 -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{- | Encode a 64-bit integer with zigzag.

In C, this is:

> (n << 1) ^ (n >> 63)
-}
toZigzag64 :: Int64 -> Word64
toZigzag64 :: Int64 -> Word64
toZigzag64 Int64
i =
  forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int64 @Word64 (Int64 -> Word64) -> Int64 -> Word64
forall a b. (a -> b) -> a -> b
$
    (Int64
i Int64 -> Int -> Int64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
1) Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
`xor` (Int64
i Int64 -> Int -> Int64
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
63)

{- | Decode a 64-bit zigzag-encoded integer.

In C, this is:

> (n >> 1) ^ (~(n & 1) + 1)
-}
fromZigzag64 :: Word64 -> Int64
fromZigzag64 :: Word64 -> Int64
fromZigzag64 Word64
n =
  Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int64) -> Word64 -> Int64
forall a b. (a -> b) -> a -> b
$
    (Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64 -> Word64
forall a. Bits a => a -> a
complement (Word64
n Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
1) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1)