{-# 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 (unsafeShiftL, unsafeShiftR, complement, (.&.), 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
n Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`mod` Natural
2 Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0 = 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 = (Integral Int64, Num Word64) => Int64 -> Word64
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)